Enhance RAG configuration and assistant loading in Neo API

- Updated the SetRAG function to include RAG settings, allowing for more detailed configuration during assistant initialization.
- Refactored the LoadBuiltIn function to streamline the saving process of assistants, improving error handling and code clarity.
- Modified loadPrompts and loadScript functions to return timestamps, enabling tracking of creation and update times for assistants.
- Enhanced the Assistant struct with created_at and updated_at fields, improving data management and traceability.
- Overall improvements to the assistant loading process, ensuring better organization and maintainability of the codebase.

neo/assistant/assistant.go

@@ -0,0 +1,253 @@
+package assistant
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/fatih/color"
+	jsoniter "github.com/json-iterator/go"
+	"github.com/yaoapp/gou/rag/driver"
+	"github.com/yaoapp/kun/log"
+)
+
+// Save save the assistant
+func (ast *Assistant) Save() error {
+	if storage == nil {
+		return fmt.Errorf("storage is not set")
+	}
+
+	_, 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
+func (ast *Assistant) Map() map[string]interface{} {
+
+	if ast == nil {
+		return nil
+	}
+
+	return map[string]interface{}{
+		"assistant_id": ast.ID,
+		"type":         ast.Type,
+		"name":         ast.Name,
+		"readonly":     ast.Readonly,
+		"avatar":       ast.Avatar,
+		"connector":    ast.Connector,
+		"path":         ast.Path,
+		"built_in":     ast.BuiltIn,
+		"sort":         ast.Sort,
+		"description":  ast.Description,
+		"options":      ast.Options,
+		"prompts":      ast.Prompts,
+		"tags":         ast.Tags,
+		"mentionable":  ast.Mentionable,
+		"automated":    ast.Automated,
+		"created_at":   timeToMySQLFormat(ast.CreatedAt),
+		"updated_at":   timeToMySQLFormat(ast.UpdatedAt),
+	}
+}
+
+// Validate validates the assistant configuration
+func (ast *Assistant) Validate() error {
+	if ast.ID == "" {
+		return fmt.Errorf("assistant_id is required")
+	}
+	if ast.Name == "" {
+		return fmt.Errorf("name is required")
+	}
+	if ast.Connector == "" {
+		return fmt.Errorf("connector is required")
+	}
+	return nil
+}
+
+// Clone creates a deep copy of the assistant
+func (ast *Assistant) Clone() *Assistant {
+	if ast == nil {
+		return nil
+	}
+
+	clone := &Assistant{
+		ID:          ast.ID,
+		Type:        ast.Type,
+		Name:        ast.Name,
+		Avatar:      ast.Avatar,
+		Connector:   ast.Connector,
+		Path:        ast.Path,
+		BuiltIn:     ast.BuiltIn,
+		Sort:        ast.Sort,
+		Description: ast.Description,
+		Readonly:    ast.Readonly,
+		Mentionable: ast.Mentionable,
+		Automated:   ast.Automated,
+		Script:      ast.Script,
+		API:         ast.API,
+	}
+
+	// Deep copy tags
+	if ast.Tags != nil {
+		clone.Tags = make([]string, len(ast.Tags))
+		copy(clone.Tags, ast.Tags)
+	}
+
+	// Deep copy options
+	if ast.Options != nil {
+		clone.Options = make(map[string]interface{})
+		for k, v := range ast.Options {
+			clone.Options[k] = v
+		}
+	}
+
+	// Deep copy prompts
+	if ast.Prompts != nil {
+		clone.Prompts = make([]Prompt, len(ast.Prompts))
+		copy(clone.Prompts, ast.Prompts)
+	}
+
+	// Deep copy flows
+	if ast.Flows != nil {
+		clone.Flows = make([]map[string]interface{}, len(ast.Flows))
+		for i, flow := range ast.Flows {
+			cloneFlow := make(map[string]interface{})
+			for k, v := range flow {
+				cloneFlow[k] = v
+			}
+			clone.Flows[i] = cloneFlow
+		}
+	}
+
+	return clone
+}
+
+// Update updates the assistant properties
+func (ast *Assistant) Update(data map[string]interface{}) error {
+	if ast == nil {
+		return fmt.Errorf("assistant is nil")
+	}
+
+	if v, ok := data["name"].(string); ok {
+		ast.Name = v
+	}
+	if v, ok := data["avatar"].(string); ok {
+		ast.Avatar = v
+	}
+	if v, ok := data["description"].(string); ok {
+		ast.Description = v
+	}
+	if v, ok := data["connector"].(string); ok {
+		ast.Connector = v
+	}
+	if v, ok := data["type"].(string); ok {
+		ast.Type = v
+	}
+	if v, ok := data["sort"].(int); ok {
+		ast.Sort = v
+	}
+	if v, ok := data["mentionable"].(bool); ok {
+		ast.Mentionable = v
+	}
+	if v, ok := data["automated"].(bool); ok {
+		ast.Automated = v
+	}
+	if v, ok := data["tags"].([]string); ok {
+		ast.Tags = v
+	}
+	if v, ok := data["options"].(map[string]interface{}); ok {
+		ast.Options = v
+	}
+
+	return ast.Validate()
+}

neo/assistant/load.go

@@ -67,16 +67,15 @@ func LoadBuiltIn() error {
 			assistant.Tags = []string{"Built-in"}
 		}
 
+		// Save the assistant
+		err = assistant.Save()
+		if err != nil {
+			return err
+		}
+
 		sort++
 		loaded.Put(assistant)
 
-		// Save the assistant
-		if storage != nil {
-			_, err := storage.SaveAssistant(assistant.Map())
-			if err != nil {
-				return err
-			}
-		}
 	}
 
 	return nil
@@ -91,11 +90,12 @@ func SetStorage(s store.Store) {
 // e: the RAG engine
 // u: the RAG file uploader
 // v: the RAG vectorizer
-func SetRAG(e driver.Engine, u driver.FileUpload, v driver.Vectorizer) {
+func SetRAG(e driver.Engine, u driver.FileUpload, v driver.Vectorizer, setting RAGSetting) {
 	rag = &RAG{
 		Engine:     e,
 		Uploader:   u,
 		Vectorizer: v,
+		Setting:    setting,
 	}
 }
 
@@ -177,24 +177,30 @@ func LoadPath(path string) (*Assistant, error) {
 	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, err := loadPrompts(promptsfile, path)
+		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, err := loadScript(scriptfile, path)
+		script, ts, err := loadScript(scriptfile, path)
 		if err != nil {
 			return nil, err
 		}
 		data["script"] = script
+		data["updated_at"] = max(updatedAt, ts)
 	}
 
 	// load functions
@@ -307,19 +313,42 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
 		}
 	}
 
+	// 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, error) {
+func loadPrompts(file string, root string) (string, int64, error) {
 
 	app, err := fs.Get("app")
 	if err != nil {
-		return "", err
+		return "", 0, err
+	}
+
+	ts, err := app.ModTime(file)
+	if err != nil {
+		return "", 0, err
 	}
 
 	prompts, err := app.ReadFile(file)
 	if err != nil {
-		return "", err
+		return "", 0, err
 	}
 
 	re := regexp.MustCompile(`@assets/([^\s]+\.(md|yml|yaml|json|txt))`)
@@ -339,11 +368,27 @@ func loadPrompts(file string, root string) (string, error) {
 		return []byte(formattedContent)
 	})
 
-	return string(prompts), nil
+	return string(prompts), ts.UnixNano(), nil
 }
 
-func loadScript(file string, root string) (*v8.Script, error) {
+func loadScript(file string, root string) (*v8.Script, int64, error) {
-	return v8.Load(file, share.ID(root, file))
+
+	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) {
@@ -353,151 +398,3 @@ func loadScriptSource(source string, file string) (*v8.Script, error) {
 	}
 	return script, nil
 }
-
-// Save save the assistant
-func (ast *Assistant) Save() error {
-	if storage == nil {
-		return fmt.Errorf("storage is not set")
-	}
-
-	_, err := storage.SaveAssistant(ast.Map())
-	return err
-}
-
-// Map convert the assistant to a map
-func (ast *Assistant) Map() map[string]interface{} {
-
-	if ast == nil {
-		return nil
-	}
-
-	return map[string]interface{}{
-		"assistant_id": ast.ID,
-		"type":         ast.Type,
-		"name":         ast.Name,
-		"readonly":     ast.Readonly,
-		"avatar":       ast.Avatar,
-		"connector":    ast.Connector,
-		"path":         ast.Path,
-		"built_in":     ast.BuiltIn,
-		"sort":         ast.Sort,
-		"description":  ast.Description,
-		"options":      ast.Options,
-		"prompts":      ast.Prompts,
-		"tags":         ast.Tags,
-		"mentionable":  ast.Mentionable,
-		"automated":    ast.Automated,
-	}
-}
-
-// Validate validates the assistant configuration
-func (ast *Assistant) Validate() error {
-	if ast.ID == "" {
-		return fmt.Errorf("assistant_id is required")
-	}
-	if ast.Name == "" {
-		return fmt.Errorf("name is required")
-	}
-	if ast.Connector == "" {
-		return fmt.Errorf("connector is required")
-	}
-	return nil
-}
-
-// Clone creates a deep copy of the assistant
-func (ast *Assistant) Clone() *Assistant {
-	if ast == nil {
-		return nil
-	}
-
-	clone := &Assistant{
-		ID:          ast.ID,
-		Type:        ast.Type,
-		Name:        ast.Name,
-		Avatar:      ast.Avatar,
-		Connector:   ast.Connector,
-		Path:        ast.Path,
-		BuiltIn:     ast.BuiltIn,
-		Sort:        ast.Sort,
-		Description: ast.Description,
-		Readonly:    ast.Readonly,
-		Mentionable: ast.Mentionable,
-		Automated:   ast.Automated,
-		Script:      ast.Script,
-		API:         ast.API,
-	}
-
-	// Deep copy tags
-	if ast.Tags != nil {
-		clone.Tags = make([]string, len(ast.Tags))
-		copy(clone.Tags, ast.Tags)
-	}
-
-	// Deep copy options
-	if ast.Options != nil {
-		clone.Options = make(map[string]interface{})
-		for k, v := range ast.Options {
-			clone.Options[k] = v
-		}
-	}
-
-	// Deep copy prompts
-	if ast.Prompts != nil {
-		clone.Prompts = make([]Prompt, len(ast.Prompts))
-		copy(clone.Prompts, ast.Prompts)
-	}
-
-	// Deep copy flows
-	if ast.Flows != nil {
-		clone.Flows = make([]map[string]interface{}, len(ast.Flows))
-		for i, flow := range ast.Flows {
-			cloneFlow := make(map[string]interface{})
-			for k, v := range flow {
-				cloneFlow[k] = v
-			}
-			clone.Flows[i] = cloneFlow
-		}
-	}
-
-	return clone
-}
-
-// Update updates the assistant properties
-func (ast *Assistant) Update(data map[string]interface{}) error {
-	if ast == nil {
-		return fmt.Errorf("assistant is nil")
-	}
-
-	if v, ok := data["name"].(string); ok {
-		ast.Name = v
-	}
-	if v, ok := data["avatar"].(string); ok {
-		ast.Avatar = v
-	}
-	if v, ok := data["description"].(string); ok {
-		ast.Description = v
-	}
-	if v, ok := data["connector"].(string); ok {
-		ast.Connector = v
-	}
-	if v, ok := data["type"].(string); ok {
-		ast.Type = v
-	}
-	if v, ok := data["sort"].(int); ok {
-		ast.Sort = v
-	}
-	if v, ok := data["mentionable"].(bool); ok {
-		ast.Mentionable = v
-	}
-	if v, ok := data["automated"].(bool); ok {
-		ast.Automated = v
-	}
-	if v, ok := data["tags"].([]string); ok {
-		ast.Tags = v
-	}
-	if v, ok := data["options"].(map[string]interface{}); ok {
-		ast.Options = v
-	}
-
-	return ast.Validate()
-}

neo/assistant/types.go

@@ -22,6 +22,12 @@ 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
@@ -59,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

@@ -119,7 +119,14 @@ func (neo *DSL) initAssistant() error {
 
 	// Assistant RAG
 	if Neo.RAG != nil {
-		assistant.SetRAG(Neo.RAG.Engine(), Neo.RAG.FileUpload(), Neo.RAG.Vectorizer())
+		assistant.SetRAG(
+			Neo.RAG.Engine(),
+			Neo.RAG.FileUpload(),
+			Neo.RAG.Vectorizer(),
+			assistant.RAGSetting{
+				IndexPrefix: Neo.RAGSetting.IndexPrefix,
+			},
+		)
 	}
 
 	// Load Built-in Assistants