diff --git a/neo/assistant/assistant.go b/neo/assistant/assistant.go index ee57d890..fd24eb9f 100644 --- a/neo/assistant/assistant.go +++ b/neo/assistant/assistant.go @@ -1,345 +1,16 @@ package assistant import ( + "context" "fmt" - "path/filepath" - "regexp" - "strings" "time" + "github.com/fatih/color" jsoniter "github.com/json-iterator/go" - "github.com/yaoapp/gou/fs" - v8 "github.com/yaoapp/gou/runtime/v8" - "github.com/yaoapp/yao/neo/store" - "github.com/yaoapp/yao/share" - "gopkg.in/yaml.v3" + "github.com/yaoapp/gou/rag/driver" + "github.com/yaoapp/kun/log" ) -// loaded the loaded assistant -var loaded = NewCache(200) // 200 is the default capacity -var storage store.Store = nil - -// LoadBuiltIn load the built-in assistants -func LoadBuiltIn() error { - root := `/assistants` - app, err := fs.Get("app") - if err != nil { - return err - } - - // Remove the built-in assistants - if storage != nil { - builtIn := true - _, err := storage.DeleteAssistants(store.AssistantFilter{BuiltIn: &builtIn}) - if err != nil { - return err - } - } - - // Check if the assistant is built-in - if exists, _ := app.Exists(root); !exists { - return nil - } - - paths, err := app.ReadDir(root, true) - if err != nil { - return err - } - - sort := 1 - for _, path := range paths { - pkgfile := filepath.Join(path, "package.yao") - if has, _ := app.Exists(pkgfile); !has { - continue - } - - assistant, err := LoadPath(path) - if err != nil { - return err - } - - assistant.Readonly = true - assistant.BuiltIn = true - assistant.Sort = sort - if assistant.Tags == nil { - assistant.Tags = []string{"Built-in"} - } - - sort++ - loaded.Put(assistant) - - // Save the assistant - if storage != nil { - _, err := storage.SaveAssistant(assistant.Map()) - if err != nil { - return err - } - } - } - - return nil -} - -// SetStorage set the storage -func SetStorage(s store.Store) { - storage = s -} - -// SetCache set the cache -func SetCache(capacity int) { - ClearCache() - loaded = NewCache(capacity) -} - -// ClearCache clear the cache -func ClearCache() { - if loaded != nil { - loaded.Clear() - loaded = nil - } -} - -// LoadStore create a new assistant from store -func LoadStore(id string) (*Assistant, error) { - assistant, exists := loaded.Get(id) - if exists { - return assistant, nil - } - - if storage == nil { - return nil, fmt.Errorf("storage is not set") - } - - data, err := storage.GetAssistant(id) - if err != nil { - return nil, err - } - - // Load from path - if data["path"] != nil { - assistant, err = LoadPath(data["path"].(string)) - if err != nil { - return nil, err - } - loaded.Put(assistant) - return assistant, nil - } - - // Load from store - assistant, err = loadMap(data) - if err != nil { - return nil, err - } - - loaded.Put(assistant) - return assistant, nil -} - -// LoadPath load assistant from path -func LoadPath(path string) (*Assistant, error) { - app, err := fs.Get("app") - if err != nil { - return nil, err - } - - pkgfile := filepath.Join(path, "package.yao") - if has, _ := app.Exists(pkgfile); !has { - return nil, fmt.Errorf("package.yao not found in %s", path) - } - - pkg, err := app.ReadFile(pkgfile) - if err != nil { - return nil, err - } - - id := strings.ReplaceAll(strings.TrimPrefix(path, "/assistants/"), "/", ".") - var data map[string]interface{} - err = jsoniter.Unmarshal(pkg, &data) - if err != nil { - return nil, err - } - - // assistant_id - data["assistant_id"] = id - data["type"] = "assistant" - data["path"] = path - // prompts - promptsfile := filepath.Join(path, "prompts.yml") - if has, _ := app.Exists(promptsfile); has { - prompts, err := loadPrompts(promptsfile, path) - if err != nil { - return nil, err - } - data["prompts"] = prompts - } - - // load script - scriptfile := filepath.Join(path, "src", "index.ts") - if has, _ := app.Exists(scriptfile); has { - script, err := loadScript(scriptfile, path) - if err != nil { - return nil, err - } - data["script"] = script - } - - // load functions - - // load flow - - return loadMap(data) -} - -func loadMap(data map[string]interface{}) (*Assistant, error) { - - assistant := &Assistant{} - - // assistant_id is required - id, ok := data["assistant_id"].(string) - if !ok { - return nil, fmt.Errorf("assistant_id is required") - } - assistant.ID = id - - // name is required - name, ok := data["name"].(string) - if !ok { - return nil, fmt.Errorf("name is required") - } - assistant.Name = name - - // avatar - if avatar, ok := data["avatar"].(string); ok { - assistant.Avatar = avatar - } - - // Type - if v, ok := data["type"].(string); ok { - assistant.Type = v - } - - // Mentionable - if v, ok := data["mentionable"].(bool); ok { - assistant.Mentionable = v - } - - // Automated - if v, ok := data["automated"].(bool); ok { - assistant.Automated = v - } - - // Readonly - if v, ok := data["readonly"].(bool); ok { - assistant.Readonly = v - } - - // built_in - if v, ok := data["built_in"].(bool); ok { - assistant.BuiltIn = v - } - - // sort - if v, ok := data["sort"].(int); ok { - assistant.Sort = v - } - - // path - if v, ok := data["path"].(string); ok { - assistant.Path = v - } - - // connector - if connector, ok := data["connector"].(string); ok { - assistant.Connector = connector - } - - // tags - if v, ok := data["tags"].([]string); ok { - assistant.Tags = v - } - - // options - if v, ok := data["options"].(map[string]interface{}); ok { - assistant.Options = v - } - - // description - if v, ok := data["description"].(string); ok { - assistant.Description = v - } - - // prompts - if v, ok := data["prompts"].(string); ok { - var prompts []Prompt - err := yaml.Unmarshal([]byte(v), &prompts) - if err != nil { - return nil, err - } - assistant.Prompts = prompts - } - - // script - if data["script"] != nil { - switch v := data["script"].(type) { - case string: - file := fmt.Sprintf("assistants/%s/src/index.ts", assistant.ID) - script, err := loadScriptSource(v, file) - if err != nil { - return nil, err - } - assistant.Script = script - case *v8.Script: - assistant.Script = v - } - } - - return assistant, nil -} - -func loadPrompts(file string, root string) (string, error) { - - app, err := fs.Get("app") - if err != nil { - return "", err - } - - prompts, err := app.ReadFile(file) - if err != nil { - return "", err - } - - re := regexp.MustCompile(`@assets/([^\s]+\.(md|yml|yaml|json|txt))`) - prompts = re.ReplaceAllFunc(prompts, func(s []byte) []byte { - asset := re.FindStringSubmatch(string(s))[1] - assetFile := filepath.Join(root, "assets", asset) - assetContent, err := app.ReadFile(assetFile) - if err != nil { - return []byte("") - } - // Add proper YAML formatting for content - lines := strings.Split(string(assetContent), "\n") - formattedContent := "|\n" - for _, line := range lines { - formattedContent += " " + line + "\n" - } - return []byte(formattedContent) - }) - - return string(prompts), nil -} - -func loadScript(file string, root string) (*v8.Script, error) { - return v8.Load(file, share.ID(root, file)) -} - -func loadScriptSource(source string, file string) (*v8.Script, error) { - script, err := v8.MakeScript([]byte(source), file, 5*time.Second, true) - if err != nil { - return nil, err - } - return script, nil -} - // Save save the assistant func (ast *Assistant) Save() error { if storage == nil { @@ -347,7 +18,98 @@ func (ast *Assistant) Save() error { } _, err := storage.SaveAssistant(ast.Map()) - return err + if err != nil { + return err + } + + // Update Index in background + go func() { + err := ast.UpdateIndex() + if err != nil { + log.Error("failed to update index for assistant %s: %s", ast.ID, err) + color.Red("failed to update index for assistant %s: %s", ast.ID, err) + } + }() + + return nil +} + +// UpdateIndex update the index for RAG +func (ast *Assistant) UpdateIndex() error { + + // RAG is not enabled + if rag == nil { + return nil + } + + if rag.Engine == nil { + return fmt.Errorf("engine is not set") + } + + // Update Index + index := fmt.Sprintf("%sassistants", rag.Setting.IndexPrefix) + id := fmt.Sprintf("assistant_%s", ast.ID) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + // Check if the index exists + exists, err := rag.Engine.HasIndex(ctx, index) + if err != nil { + return err + } + + // Create the index if it does not exist + if !exists { + ctxCreate, cancelCreate := context.WithTimeout(context.Background(), 2*time.Second) + defer cancelCreate() + err = rag.Engine.CreateIndex(ctxCreate, driver.IndexConfig{Name: index}) + if err != nil { + return err + } + } + + // Check if the document exists + exists, err = rag.Engine.HasDocument(ctx, index, id) + if err != nil { + return err + } + + // Check if the document is updated + if exists { + metadata, err := rag.Engine.GetMetadata(ctx, index, id) + if err != nil { + return err + } + + if v, ok := metadata["updated_at"].(string); ok { + updatedAt, err := stringToTimestamp(v) + if err != nil { + return err + } + if updatedAt >= ast.UpdatedAt { + return nil + } + } + } + + // Update the index + content, err := jsoniter.MarshalToString(ast.Map()) + if err != nil { + return err + } + + metadata := map[string]interface{}{ + "assistant_id": ast.ID, + "type": ast.Type, + "name": ast.Name, + "updated_at": fmt.Sprintf("%d", ast.UpdatedAt), + } + + return rag.Engine.IndexDoc(ctx, index, &driver.Document{ + DocID: id, + Content: content, + Metadata: metadata, + }) } // Map convert the assistant to a map @@ -373,6 +135,8 @@ func (ast *Assistant) Map() map[string]interface{} { "tags": ast.Tags, "mentionable": ast.Mentionable, "automated": ast.Automated, + "created_at": timeToMySQLFormat(ast.CreatedAt), + "updated_at": timeToMySQLFormat(ast.UpdatedAt), } } diff --git a/neo/assistant/load.go b/neo/assistant/load.go new file mode 100644 index 00000000..734d612e --- /dev/null +++ b/neo/assistant/load.go @@ -0,0 +1,400 @@ +package assistant + +import ( + "fmt" + "path/filepath" + "regexp" + "strings" + "time" + + 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" + "gopkg.in/yaml.v3" +) + +// 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 { + root := `/assistants` + app, err := fs.Get("app") + if err != nil { + return err + } + + // Remove the built-in assistants + if storage != nil { + builtIn := true + _, err := storage.DeleteAssistants(store.AssistantFilter{BuiltIn: &builtIn}) + if err != nil { + return err + } + } + + // Check if the assistant is built-in + if exists, _ := app.Exists(root); !exists { + return nil + } + + paths, err := app.ReadDir(root, true) + if err != nil { + return err + } + + sort := 1 + for _, path := range paths { + pkgfile := filepath.Join(path, "package.yao") + if has, _ := app.Exists(pkgfile); !has { + continue + } + + assistant, err := LoadPath(path) + if err != nil { + return err + } + + assistant.Readonly = true + assistant.BuiltIn = true + assistant.Sort = sort + if assistant.Tags == nil { + assistant.Tags = []string{"Built-in"} + } + + // Save the assistant + err = assistant.Save() + if err != nil { + return err + } + + sort++ + loaded.Put(assistant) + + } + + return nil +} + +// SetStorage set the storage +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, setting RAGSetting) { + rag = &RAG{ + Engine: e, + Uploader: u, + Vectorizer: v, + Setting: setting, + } +} + +// SetCache set the cache +func SetCache(capacity int) { + ClearCache() + loaded = NewCache(capacity) +} + +// ClearCache clear the cache +func ClearCache() { + if loaded != nil { + loaded.Clear() + loaded = nil + } +} + +// LoadStore create a new assistant from store +func LoadStore(id string) (*Assistant, error) { + assistant, exists := loaded.Get(id) + if exists { + return assistant, nil + } + + if storage == nil { + return nil, fmt.Errorf("storage is not set") + } + + data, err := storage.GetAssistant(id) + if err != nil { + return nil, err + } + + // Load from path + if data["path"] != nil { + assistant, err = LoadPath(data["path"].(string)) + if err != nil { + return nil, err + } + loaded.Put(assistant) + return assistant, nil + } + + // Load from store + assistant, err = loadMap(data) + if err != nil { + return nil, err + } + + loaded.Put(assistant) + return assistant, nil +} + +// LoadPath load assistant from path +func LoadPath(path string) (*Assistant, error) { + app, err := fs.Get("app") + if err != nil { + return nil, err + } + + pkgfile := filepath.Join(path, "package.yao") + if has, _ := app.Exists(pkgfile); !has { + return nil, fmt.Errorf("package.yao not found in %s", path) + } + + pkg, err := app.ReadFile(pkgfile) + if err != nil { + return nil, err + } + + id := strings.ReplaceAll(strings.TrimPrefix(path, "/assistants/"), "/", ".") + var data map[string]interface{} + err = jsoniter.Unmarshal(pkg, &data) + if err != nil { + return nil, err + } + + // assistant_id + data["assistant_id"] = id + data["type"] = "assistant" + data["path"] = path + + updatedAt := int64(0) + + // prompts + promptsfile := filepath.Join(path, "prompts.yml") + if has, _ := app.Exists(promptsfile); has { + prompts, ts, err := loadPrompts(promptsfile, path) + if err != nil { + return nil, err + } + data["prompts"] = prompts + data["updated_at"] = ts + updatedAt = ts + } + + // load script + scriptfile := filepath.Join(path, "src", "index.ts") + if has, _ := app.Exists(scriptfile); has { + script, ts, err := loadScript(scriptfile, path) + if err != nil { + return nil, err + } + data["script"] = script + data["updated_at"] = max(updatedAt, ts) + } + + // load functions + + // load flow + + return loadMap(data) +} + +func loadMap(data map[string]interface{}) (*Assistant, error) { + + assistant := &Assistant{} + + // assistant_id is required + id, ok := data["assistant_id"].(string) + if !ok { + return nil, fmt.Errorf("assistant_id is required") + } + assistant.ID = id + + // name is required + name, ok := data["name"].(string) + if !ok { + return nil, fmt.Errorf("name is required") + } + assistant.Name = name + + // avatar + if avatar, ok := data["avatar"].(string); ok { + assistant.Avatar = avatar + } + + // Type + if v, ok := data["type"].(string); ok { + assistant.Type = v + } + + // Mentionable + if v, ok := data["mentionable"].(bool); ok { + assistant.Mentionable = v + } + + // Automated + if v, ok := data["automated"].(bool); ok { + assistant.Automated = v + } + + // Readonly + if v, ok := data["readonly"].(bool); ok { + assistant.Readonly = v + } + + // built_in + if v, ok := data["built_in"].(bool); ok { + assistant.BuiltIn = v + } + + // sort + if v, ok := data["sort"].(int); ok { + assistant.Sort = v + } + + // path + if v, ok := data["path"].(string); ok { + assistant.Path = v + } + + // connector + if connector, ok := data["connector"].(string); ok { + assistant.Connector = connector + } + + // tags + if v, ok := data["tags"].([]string); ok { + assistant.Tags = v + } + + // options + if v, ok := data["options"].(map[string]interface{}); ok { + assistant.Options = v + } + + // description + if v, ok := data["description"].(string); ok { + assistant.Description = v + } + + // prompts + if v, ok := data["prompts"].(string); ok { + var prompts []Prompt + err := yaml.Unmarshal([]byte(v), &prompts) + if err != nil { + return nil, err + } + assistant.Prompts = prompts + } + + // script + if data["script"] != nil { + switch v := data["script"].(type) { + case string: + file := fmt.Sprintf("assistants/%s/src/index.ts", assistant.ID) + script, err := loadScriptSource(v, file) + if err != nil { + return nil, err + } + assistant.Script = script + case *v8.Script: + assistant.Script = v + } + } + + // created_at + if v, has := data["created_at"]; has { + ts, err := getTimestamp(v) + if err != nil { + return nil, err + } + assistant.CreatedAt = ts + } + + // updated_at + if v, has := data["updated_at"]; has { + ts, err := getTimestamp(v) + if err != nil { + return nil, err + } + assistant.UpdatedAt = ts + } + + return assistant, nil +} + +func loadPrompts(file string, root string) (string, int64, error) { + + app, err := fs.Get("app") + if err != nil { + return "", 0, err + } + + ts, err := app.ModTime(file) + if err != nil { + return "", 0, err + } + + prompts, err := app.ReadFile(file) + if err != nil { + return "", 0, err + } + + re := regexp.MustCompile(`@assets/([^\s]+\.(md|yml|yaml|json|txt))`) + prompts = re.ReplaceAllFunc(prompts, func(s []byte) []byte { + asset := re.FindStringSubmatch(string(s))[1] + assetFile := filepath.Join(root, "assets", asset) + assetContent, err := app.ReadFile(assetFile) + if err != nil { + return []byte("") + } + // Add proper YAML formatting for content + lines := strings.Split(string(assetContent), "\n") + formattedContent := "|\n" + for _, line := range lines { + formattedContent += " " + line + "\n" + } + return []byte(formattedContent) + }) + + return string(prompts), ts.UnixNano(), nil +} + +func loadScript(file string, root string) (*v8.Script, int64, error) { + + app, err := fs.Get("app") + if err != nil { + return nil, 0, err + } + + ts, err := app.ModTime(file) + if err != nil { + return nil, 0, err + } + + script, err := v8.Load(file, share.ID(root, file)) + if err != nil { + return nil, 0, err + } + + return script, ts.UnixNano(), nil +} + +func loadScriptSource(source string, file string) (*v8.Script, error) { + script, err := v8.MakeScript([]byte(source), file, 5*time.Second, true) + if err != nil { + return nil, err + } + return script, nil +} diff --git a/neo/assistant/assistant_test.go b/neo/assistant/load_test.go similarity index 97% rename from neo/assistant/assistant_test.go rename to neo/assistant/load_test.go index fb4f20e1..5c6d21ba 100644 --- a/neo/assistant/assistant_test.go +++ b/neo/assistant/load_test.go @@ -14,7 +14,7 @@ func prepare(t *testing.T) { test.Prepare(t, config.Conf) } -func TestAssistant_LoadPath(t *testing.T) { +func TestLoad_LoadPath(t *testing.T) { prepare(t) defer test.Clean() @@ -37,7 +37,7 @@ func TestAssistant_LoadPath(t *testing.T) { assert.Error(t, err) } -func TestAssistant_LoadStore(t *testing.T) { +func TestLoad_LoadStore(t *testing.T) { prepare(t) defer test.Clean() @@ -79,7 +79,7 @@ func TestAssistant_LoadStore(t *testing.T) { assert.Error(t, err) } -func TestAssistant_Cache(t *testing.T) { +func TestLoad_Cache(t *testing.T) { prepare(t) defer test.Clean() @@ -127,7 +127,7 @@ func TestAssistant_Cache(t *testing.T) { assert.NotNil(t, loaded) } -func TestAssistant_Validate(t *testing.T) { +func TestLoad_Validate(t *testing.T) { tests := []struct { name string ast *Assistant @@ -178,7 +178,7 @@ func TestAssistant_Validate(t *testing.T) { } } -func TestAssistant_Clone(t *testing.T) { +func TestLoad_Clone(t *testing.T) { // Create a test assistant with all fields populated original := &Assistant{ ID: "test-id", @@ -233,7 +233,7 @@ func TestAssistant_Clone(t *testing.T) { assert.Nil(t, nilAssistant.Clone()) } -func TestAssistant_Update(t *testing.T) { +func TestLoad_Update(t *testing.T) { // Create a test assistant ast := &Assistant{ ID: "test-id", diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 977cf5e9..5cec9ba6 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,19 @@ 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 + Setting RAGSetting +} + +// RAGSetting the RAG setting +type RAGSetting struct { + IndexPrefix string `json:"index_prefix" yaml:"index_prefix"` +} + // Prompt a prompt type Prompt struct { Role string `json:"role"` @@ -51,6 +65,8 @@ type Assistant struct { Flows []map[string]interface{} `json:"flows,omitempty"` // Assistant Flows Script *v8.Script `json:"-" yaml:"-"` // Assistant Script API API `json:"-" yaml:"-"` // Assistant API + CreatedAt int64 `json:"created_at"` // Creation timestamp + UpdatedAt int64 `json:"updated_at"` // Last update timestamp } // File the file diff --git a/neo/assistant/utils.go b/neo/assistant/utils.go new file mode 100644 index 00000000..aa7c6720 --- /dev/null +++ b/neo/assistant/utils.go @@ -0,0 +1,44 @@ +package assistant + +import ( + "fmt" + "strconv" + "time" +) + +func getTimestamp(v interface{}) (int64, error) { + switch v := v.(type) { + case int64: + return v, nil + case int: + return int64(v), nil + + case string: + if ts, err := time.Parse(time.RFC3339, v); err == nil { + return ts.UnixNano(), nil + } + + // MySQL format + if ts, err := time.Parse("2006-01-02 15:04:05", v); err == nil { + return ts.UnixNano(), nil + } + + // UnixNano format + if ts, err := strconv.ParseInt(v, 10, 64); err == nil { + return ts, nil + } + + } + return 0, fmt.Errorf("invalid timestamp type") +} + +func stringToTimestamp(v string) (int64, error) { + return strconv.ParseInt(v, 10, 64) +} + +func timeToMySQLFormat(ts int64) string { + if ts == 0 { + return "0000-00-00 00:00:00" + } + return time.Unix(ts/1e9, ts%1e9).Format("2006-01-02 15:04:05") +} diff --git a/neo/load.go b/neo/load.go index ef45f86a..2e364b6c 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 { @@ -38,7 +26,7 @@ func Load(cfg config.Config) error { Option: map[string]interface{}{}, Allows: []string{}, StoreSetting: store.Setting{ - Table: "yao_neo_conversation", + Prefix: "yao_neo_", Connector: "default", }, } @@ -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,86 @@ 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(), + assistant.RAGSetting{ + IndexPrefix: Neo.RAGSetting.IndexPrefix, + }, + ) + } + + // 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 +144,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() diff --git a/neo/store/types.go b/neo/store/types.go index 925c5ecd..1d6a889b 100644 --- a/neo/store/types.go +++ b/neo/store/types.go @@ -5,7 +5,7 @@ package store type Setting struct { Connector string `json:"connector,omitempty"` // Name of the connector used to specify data storage method UserField string `json:"user_field,omitempty"` // User ID field name, defaults to "user_id" - Table string `json:"table,omitempty"` // Database table name + Prefix string `json:"prefix,omitempty"` // Database table name prefix MaxSize int `json:"max_size,omitempty" yaml:"max_size,omitempty"` // Maximum storage size limit TTL int `json:"ttl,omitempty" yaml:"ttl,omitempty"` // Time To Live in seconds } diff --git a/neo/store/xun.go b/neo/store/xun.go index 05c9c068..aadb4583 100644 --- a/neo/store/xun.go +++ b/neo/store/xun.go @@ -99,7 +99,7 @@ func (conv *Xun) clean() { } if nums > 0 { - log.Trace("Clean the conversation table: %s %d", conv.setting.Table, nums) + log.Trace("Clean the conversation table: %s %d", conv.setting.Prefix, nums) } } @@ -283,15 +283,15 @@ func (conv *Xun) getUserID(sid string) (string, error) { } func (conv *Xun) getHistoryTable() string { - return conv.setting.Table + "_history" + return conv.setting.Prefix + "history" } func (conv *Xun) getChatTable() string { - return conv.setting.Table + "_chat" + return conv.setting.Prefix + "chat" } func (conv *Xun) getAssistantTable() string { - return conv.setting.Table + "_assistant" + return conv.setting.Prefix + "assistant" } // UpdateChatTitle update the chat title diff --git a/neo/store/xun_test.go b/neo/store/xun_test.go index dc42de04..4dbd936a 100644 --- a/neo/store/xun_test.go +++ b/neo/store/xun_test.go @@ -40,7 +40,7 @@ func TestNewXunDefault(t *testing.T) { store, err := NewXun(Setting{ Connector: "default", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", }) if err != nil { @@ -139,7 +139,7 @@ func TestNewXunConnector(t *testing.T) { store, err := NewXun(Setting{ Connector: "mysql", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", }) if err != nil { @@ -201,7 +201,7 @@ func TestXunSaveAndGetHistory(t *testing.T) { store, err := NewXun(Setting{ Connector: "default", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", TTL: 3600, }) @@ -239,7 +239,7 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) { store, err := NewXun(Setting{ Connector: "default", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", TTL: 3600, }) @@ -308,7 +308,7 @@ func TestXunGetChats(t *testing.T) { store, err := NewXun(Setting{ Connector: "default", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", }) if err != nil { t.Fatal(err) @@ -364,7 +364,7 @@ func TestXunDeleteChat(t *testing.T) { store, err := NewXun(Setting{ Connector: "default", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", }) if err != nil { t.Fatal(err) @@ -404,7 +404,7 @@ func TestXunDeleteAllChats(t *testing.T) { store, err := NewXun(Setting{ Connector: "default", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", }) if err != nil { t.Fatal(err) @@ -455,7 +455,7 @@ func TestXunAssistantCRUD(t *testing.T) { store, err := NewXun(Setting{ Connector: "default", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", }) if err != nil { t.Fatal(err) @@ -926,7 +926,7 @@ func TestXunAssistantPagination(t *testing.T) { store, err := NewXun(Setting{ Connector: "default", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", }) if err != nil { t.Fatal(err) @@ -1250,7 +1250,7 @@ func TestGetAssistantTags(t *testing.T) { store, err := NewXun(Setting{ Connector: "default", - Table: "__unit_test_conversation", + Prefix: "__unit_test_conversation_", }) if err != nil { t.Fatal(err)