diff --git a/neo/assistant/assistant.go b/neo/assistant/assistant.go new file mode 100644 index 000000000..fd24eb9f4 --- /dev/null +++ b/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() +} diff --git a/neo/assistant/load.go b/neo/assistant/load.go index 2460a44df..734d612e0 100644 --- a/neo/assistant/load.go +++ b/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,165 +368,33 @@ 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) { - return v8.Load(file, share.ID(root, file)) -} +func loadScript(file string, root string) (*v8.Script, int64, error) { -func loadScriptSource(source string, file string) (*v8.Script, error) { - script, err := v8.MakeScript([]byte(source), file, 5*time.Second, true) + app, err := fs.Get("app") if err != nil { - return nil, err + return nil, 0, err } - 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) + ts, err := app.ModTime(file) + if err != nil { + return nil, 0, err } - // 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 - } + script, err := v8.Load(file, share.ID(root, file)) + if err != nil { + return nil, 0, err } - return clone + return script, ts.UnixNano(), nil } -// 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 +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 ast.Validate() + return script, nil } diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 17f17bf00..5cec9ba61 100644 --- a/neo/assistant/types.go +++ b/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 diff --git a/neo/assistant/utils.go b/neo/assistant/utils.go new file mode 100644 index 000000000..aa7c6720e --- /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 cfafc65c8..2e364b6c7 100644 --- a/neo/load.go +++ b/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