Merge pull request #814 from trheyi/main

Neo API and Conversation Management Refactoring

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"
+	"github.com/yaoapp/gou/rag/driver"
-	v8 "github.com/yaoapp/gou/runtime/v8"
+	"github.com/yaoapp/kun/log"
-	"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
-
-// 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())
+	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),
 	}
 }
 

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

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

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

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

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
+	// Initialize Assistant
-	assistant.SetStorage(Neo.Store)
+	err = Neo.initAssistant()
-	err = assistant.LoadBuiltIn()
 	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)
+}

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

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
 }

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

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)