From f8ce8fc1932ab02a147aa69f24dd05d0f42e1b38 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 1 Jan 2025 16:48:06 +0800 Subject: [PATCH] Enhance assistant management and API functionality in Neo - Added a new endpoint to retrieve all assistant tags, improving data accessibility for clients. - Updated the assistant list handling to support filtering by built-in status and assistant ID, enhancing the filtering capabilities. - Introduced a method to load built-in assistants, streamlining the assistant initialization process. - Enhanced the Assistant struct with new fields for path, built-in status, and sorting, improving data organization. - Implemented validation and cloning methods for the Assistant struct, ensuring data integrity and ease of use. - Updated tests to cover new functionalities, including validation, cloning, and tag retrieval, ensuring robust functionality across the assistant management operations. --- neo/api.go | 36 +++ neo/assistant/assistant.go | 243 ++++++++++++++++- neo/assistant/assistant_test.go | 194 +++++++++++++- neo/assistant/types.go | 9 +- neo/load.go | 7 + neo/store/mongo.go | 10 + neo/store/redis.go | 10 + neo/store/types.go | 10 + neo/store/xun.go | 97 ++++++- neo/store/xun_test.go | 445 ++++++++++++++++++++++++++++++-- 10 files changed, 1039 insertions(+), 22 deletions(-) diff --git a/neo/api.go b/neo/api.go index 1ecc7dd3..aebe7bad 100644 --- a/neo/api.go +++ b/neo/api.go @@ -62,6 +62,9 @@ func (neo *DSL) API(router *gin.Engine, path string) error { // List assistants example: // curl -X GET 'http://localhost:5099/api/__yao/neo/assistants?page=1&pagesize=20&tags=tag1,tag2&token=xxx' router.GET(path+"/assistants", append(middlewares, neo.handleAssistantList)...) + // Get all assistant tags example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/assistants/tags?token=xxx' + router.GET(path+"/assistants/tags", append(middlewares, neo.handleAssistantTags)...) // Get assistant details example: // curl -X GET 'http://localhost:5099/api/__yao/neo/assistants/assistant_123?token=xxx' @@ -878,6 +881,14 @@ func (neo *DSL) handleAssistantList(c *gin.Context) { filter.Select = strings.Split(selectFields, ",") } + // Parse built_in (support various boolean formats) + if builtIn := c.Query("built_in"); builtIn != "" { + val := parseBoolValue(builtIn) + if val != nil { + filter.BuiltIn = val + } + } + // Parse mentionable (support various boolean formats) if mentionable := c.Query("mentionable"); mentionable != "" { val := parseBoolValue(mentionable) @@ -894,6 +905,11 @@ func (neo *DSL) handleAssistantList(c *gin.Context) { } } + // Parse assistant_id + if assistantID := c.Query("assistant_id"); assistantID != "" { + filter.AssistantID = assistantID + } + response, err := neo.Store.GetAssistants(filter) if err != nil { c.JSON(500, gin.H{"message": err.Error(), "code": 500}) @@ -1023,3 +1039,23 @@ func (neo *DSL) handleConnectors(c *gin.Context) { c.JSON(200, gin.H{"data": options}) c.Done() } + +// handleAssistantTags handles getting all assistant tags +func (neo *DSL) handleAssistantTags(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + tags, err := neo.Store.GetAssistantTags() + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, gin.H{"data": tags}) + c.Done() +} diff --git a/neo/assistant/assistant.go b/neo/assistant/assistant.go index f3ad5640..1a85867a 100644 --- a/neo/assistant/assistant.go +++ b/neo/assistant/assistant.go @@ -19,6 +19,63 @@ import ( 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 + 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 @@ -89,7 +146,8 @@ func LoadPath(path string) (*Assistant, error) { // 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 { @@ -140,6 +198,41 @@ func loadMap(data map[string]interface{}) (*Assistant, error) { 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 @@ -216,3 +309,151 @@ 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() +} diff --git a/neo/assistant/assistant_test.go b/neo/assistant/assistant_test.go index f3319c1c..fb4f20e1 100644 --- a/neo/assistant/assistant_test.go +++ b/neo/assistant/assistant_test.go @@ -127,6 +127,196 @@ func TestAssistant_Cache(t *testing.T) { assert.NotNil(t, loaded) } +func TestAssistant_Validate(t *testing.T) { + tests := []struct { + name string + ast *Assistant + wantErr bool + }{ + { + name: "valid assistant", + ast: &Assistant{ + ID: "test-id", + Name: "Test Assistant", + Connector: "test-connector", + }, + wantErr: false, + }, + { + name: "missing id", + ast: &Assistant{ + Name: "Test Assistant", + Connector: "test-connector", + }, + wantErr: true, + }, + { + name: "missing name", + ast: &Assistant{ + ID: "test-id", + Connector: "test-connector", + }, + wantErr: true, + }, + { + name: "missing connector", + ast: &Assistant{ + ID: "test-id", + Name: "Test Assistant", + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.ast.Validate() + if (err != nil) != tt.wantErr { + t.Errorf("Assistant.Validate() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func TestAssistant_Clone(t *testing.T) { + // Create a test assistant with all fields populated + original := &Assistant{ + ID: "test-id", + Type: "test-type", + Name: "Test Assistant", + Avatar: "test-avatar", + Connector: "test-connector", + Path: "test-path", + BuiltIn: true, + Sort: 1, + Description: "test description", + Tags: []string{"tag1", "tag2"}, + Readonly: true, + Mentionable: true, + Automated: true, + Options: map[string]interface{}{"key": "value"}, + Prompts: []Prompt{{Role: "system", Content: "test"}}, + Flows: []map[string]interface{}{{"step": "test"}}, + } + + // Clone the assistant + clone := original.Clone() + + // Verify all fields are correctly cloned + assert.Equal(t, original.ID, clone.ID) + assert.Equal(t, original.Type, clone.Type) + assert.Equal(t, original.Name, clone.Name) + assert.Equal(t, original.Avatar, clone.Avatar) + assert.Equal(t, original.Connector, clone.Connector) + assert.Equal(t, original.Path, clone.Path) + assert.Equal(t, original.BuiltIn, clone.BuiltIn) + assert.Equal(t, original.Sort, clone.Sort) + assert.Equal(t, original.Description, clone.Description) + assert.Equal(t, original.Tags, clone.Tags) + assert.Equal(t, original.Readonly, clone.Readonly) + assert.Equal(t, original.Mentionable, clone.Mentionable) + assert.Equal(t, original.Automated, clone.Automated) + assert.Equal(t, original.Options, clone.Options) + assert.Equal(t, original.Prompts, clone.Prompts) + assert.Equal(t, original.Flows, clone.Flows) + + // Verify deep copy by modifying original + original.Tags[0] = "modified" + original.Options["key"] = "modified" + original.Flows[0]["step"] = "modified" + assert.NotEqual(t, original.Tags[0], clone.Tags[0]) + assert.NotEqual(t, original.Options["key"], clone.Options["key"]) + assert.NotEqual(t, original.Flows[0]["step"], clone.Flows[0]["step"]) + + // Test nil case + var nilAssistant *Assistant + assert.Nil(t, nilAssistant.Clone()) +} + +func TestAssistant_Update(t *testing.T) { + // Create a test assistant + ast := &Assistant{ + ID: "test-id", + Name: "Original Name", + Connector: "original-connector", + } + + // Test updating various fields + updates := map[string]interface{}{ + "name": "Updated Name", + "avatar": "updated-avatar", + "description": "Updated description", + "connector": "updated-connector", + "type": "updated-type", + "sort": 2, + "mentionable": true, + "automated": true, + "tags": []string{"new-tag"}, + "options": map[string]interface{}{"new": "value"}, + } + + err := ast.Update(updates) + assert.NoError(t, err) + + // Verify updates + assert.Equal(t, "Updated Name", ast.Name) + assert.Equal(t, "updated-avatar", ast.Avatar) + assert.Equal(t, "Updated description", ast.Description) + assert.Equal(t, "updated-connector", ast.Connector) + assert.Equal(t, "updated-type", ast.Type) + assert.Equal(t, 2, ast.Sort) + assert.True(t, ast.Mentionable) + assert.True(t, ast.Automated) + assert.Equal(t, []string{"new-tag"}, ast.Tags) + assert.Equal(t, map[string]interface{}{"new": "value"}, ast.Options) + + // Test nil assistant + var nilAssistant *Assistant + err = nilAssistant.Update(updates) + assert.Error(t, err) + + // Test invalid update that would make the assistant invalid + invalidUpdates := map[string]interface{}{ + "name": "", + } + err = ast.Update(invalidUpdates) + assert.Error(t, err) +} + +func TestLoadBuiltIn(t *testing.T) { + prepare(t) + defer test.Clean() + + // Clear any existing cache and storage + ClearCache() + SetStorage(nil) + + // Create a mock store to verify built-in assistants are saved + mockStore := &mockStore{ + data: make(map[string]map[string]interface{}), + } + SetStorage(mockStore) + SetCache(100) + + // Test loading built-in assistants + err := LoadBuiltIn() + assert.NoError(t, err) + + // Verify Modi assistant was loaded + assistant, exists := loaded.Get("modi") + assert.True(t, exists, "Modi assistant should be loaded in cache") + if exists { + assert.Equal(t, "modi", assistant.ID) + assert.Equal(t, "Modi", assistant.Name) + assert.Equal(t, "deepseek", assistant.Connector) + assert.True(t, assistant.BuiltIn) + assert.True(t, assistant.Readonly) + assert.NotNil(t, assistant.Prompts) + assert.NotNil(t, assistant.Script) + } + +} + // mockStore implements store.Store interface for testing type mockStore struct { data map[string]map[string]interface{} @@ -193,4 +383,6 @@ func (m *mockStore) SaveAssistant(assistant map[string]interface{}) (interface{} func (m *mockStore) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { return nil } -func (m *mockStore) UpdateChatTitle(sid string, cid string, title string) error { return nil } +func (m *mockStore) UpdateChatTitle(sid string, cid string, title string) error { return nil } +func (m *mockStore) DeleteAssistants(filter store.AssistantFilter) (int64, error) { return 0, nil } +func (m *mockStore) GetAssistantTags() ([]string, error) { return []string{}, nil } diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 30b8c801..977cf5e9 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -38,8 +38,15 @@ type Assistant struct { Name string `json:"name,omitempty"` // Assistant Name Avatar string `json:"avatar,omitempty"` // Assistant Avatar Connector string `json:"connector"` // AI Connector + Path string `json:"path,omitempty"` // Assistant Path + BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant + Sort int `json:"sort,omitempty"` // Assistant Sort Description string `json:"description,omitempty"` // Assistant Description - Option map[string]interface{} `json:"option,omitempty"` // AI Option + Tags []string `json:"tags,omitempty"` // Assistant Tags + Readonly bool `json:"readonly,omitempty"` // Whether this assistant is readonly + Mentionable bool `json:"mentionable,omitempty"` // Whether this assistant is mentionable + Automated bool `json:"automated,omitempty"` // Whether this assistant is automated + Options map[string]interface{} `json:"options,omitempty"` // AI Options Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts Flows []map[string]interface{} `json:"flows,omitempty"` // Assistant Flows Script *v8.Script `json:"-" yaml:"-"` // Assistant Script diff --git a/neo/load.go b/neo/load.go index 0a918d81..6b87957f 100644 --- a/neo/load.go +++ b/neo/load.go @@ -51,6 +51,13 @@ func Load(cfg config.Config) error { return err } + // Load Built-in Assistants + assistant.SetStorage(Neo.Store) + err = assistant.LoadBuiltIn() + if err != nil { + return err + } + // Query Assistant List ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() diff --git a/neo/store/mongo.go b/neo/store/mongo.go index 9c7ec275..e94b4278 100644 --- a/neo/store/mongo.go +++ b/neo/store/mongo.go @@ -62,3 +62,13 @@ func (m *Mongo) GetAssistants(filter AssistantFilter) (*AssistantResponse, error func (m *Mongo) GetAssistant(assistantID string) (map[string]interface{}, error) { return map[string]interface{}{}, nil } + +// DeleteAssistants deletes assistants based on filter conditions (not implemented) +func (mongo *Mongo) DeleteAssistants(filter AssistantFilter) (int64, error) { + return 0, nil +} + +// GetAssistantTags retrieves all unique tags from assistants +func (conv *Mongo) GetAssistantTags() ([]string, error) { + return []string{}, nil +} diff --git a/neo/store/redis.go b/neo/store/redis.go index 94ab30db..50206e1a 100644 --- a/neo/store/redis.go +++ b/neo/store/redis.go @@ -62,3 +62,13 @@ func (r *Redis) GetAssistants(filter AssistantFilter) (*AssistantResponse, error func (r *Redis) GetAssistant(assistantID string) (map[string]interface{}, error) { return map[string]interface{}{}, nil } + +// DeleteAssistants deletes assistants based on filter conditions (not implemented) +func (redis *Redis) DeleteAssistants(filter AssistantFilter) (int64, error) { + return 0, nil +} + +// GetAssistantTags retrieves all unique tags from assistants +func (conv *Redis) GetAssistantTags() ([]string, error) { + return []string{}, nil +} diff --git a/neo/store/types.go b/neo/store/types.go index 628a7627..925c5ecd 100644 --- a/neo/store/types.go +++ b/neo/store/types.go @@ -52,6 +52,7 @@ type AssistantFilter struct { AssistantID string `json:"assistant_id,omitempty"` // Filter by assistant ID Mentionable *bool `json:"mentionable,omitempty"` // Filter by mentionable status Automated *bool `json:"automated,omitempty"` // Filter by automation status + BuiltIn *bool `json:"built_in,omitempty"` // Filter by built-in status Page int `json:"page,omitempty"` // Page number, starting from 1 PageSize int `json:"pagesize,omitempty"` // Items per page Select []string `json:"select,omitempty"` // Fields to return, returns all fields if empty @@ -135,4 +136,13 @@ type Store interface { // assistantID: Assistant ID // Returns: Assistant information and potential error GetAssistant(assistantID string) (map[string]interface{}, error) + + // DeleteAssistants deletes assistants based on filter conditions + // filter: Filter conditions + // Returns: Number of deleted records and potential error + DeleteAssistants(filter AssistantFilter) (int64, error) + + // GetAssistantTags retrieves all unique tags from assistants + // Returns: List of tags and potential error + GetAssistantTags() ([]string, error) } diff --git a/neo/store/xun.go b/neo/store/xun.go index 41ee1ade..05c9c068 100644 --- a/neo/store/xun.go +++ b/neo/store/xun.go @@ -225,6 +225,9 @@ func (conv *Xun) initAssistantTable() error { table.String("avatar", 200).Null() // assistant avatar table.String("connector", 200).NotNull() // assistant connector table.Text("description").Null() // assistant description + table.String("path", 200).Null() // assistant storage path + table.Integer("sort").SetDefault(9999).Index() // assistant sort order + table.Boolean("built_in").SetDefault(false).Index() // whether this is a built-in assistant table.JSON("options").Null() // assistant options table.JSON("prompts").Null() // assistant prompts table.JSON("flows").Null() // assistant flows @@ -251,7 +254,7 @@ func (conv *Xun) initAssistantTable() error { return err } - fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "tags", "mentionable", "created_at", "updated_at"} + fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "path", "sort", "built_in", "options", "prompts", "flows", "files", "functions", "tags", "mentionable", "created_at", "updated_at"} for _, field := range fields { if !tab.HasColumn(field) { return fmt.Errorf("%s is required", field) @@ -845,6 +848,11 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro qb.Where("automated", *filter.Automated) } + // Apply built_in filter if provided + if filter.BuiltIn != nil { + qb.Where("built_in", *filter.BuiltIn) + } + // Set defaults for pagination if filter.PageSize <= 0 { filter.PageSize = 20 @@ -881,7 +889,8 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro } // Get paginated results - rows, err := qb.OrderBy("created_at", "desc"). + rows, err := qb.OrderBy("sort", "asc"). + OrderBy("updated_at", "desc"). Offset(offset). Limit(filter.PageSize). Get() @@ -950,3 +959,87 @@ func (conv *Xun) GetAssistant(assistantID string) (map[string]interface{}, error return data, nil } + +// DeleteAssistants deletes assistants based on filter conditions +func (conv *Xun) DeleteAssistants(filter AssistantFilter) (int64, error) { + qb := conv.query.New(). + Table(conv.getAssistantTable()) + + // Apply tag filter if provided + if filter.Tags != nil && len(filter.Tags) > 0 { + qb.Where(func(qb query.Query) { + for i, tag := range filter.Tags { + pattern := fmt.Sprintf("%%\"%s\"%%", tag) + if i == 0 { + qb.Where("tags", "like", pattern) + } else { + qb.OrWhere("tags", "like", pattern) + } + } + }) + } + + // Apply keyword filter if provided + if filter.Keywords != "" { + qb.Where(func(qb query.Query) { + qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Keywords)). + OrWhere("description", "like", fmt.Sprintf("%%%s%%", filter.Keywords)) + }) + } + + // Apply connector filter if provided + if filter.Connector != "" { + qb.Where("connector", filter.Connector) + } + + // Apply assistant_id filter if provided + if filter.AssistantID != "" { + qb.Where("assistant_id", filter.AssistantID) + } + + // Apply mentionable filter if provided + if filter.Mentionable != nil { + qb.Where("mentionable", *filter.Mentionable) + } + + // Apply automated filter if provided + if filter.Automated != nil { + qb.Where("automated", *filter.Automated) + } + + // Apply built_in filter if provided + if filter.BuiltIn != nil { + qb.Where("built_in", *filter.BuiltIn) + } + + // Execute delete and return number of deleted records + return qb.Delete() +} + +// GetAssistantTags retrieves all unique tags from assistants +func (conv *Xun) GetAssistantTags() ([]string, error) { + q := conv.newQuery().Table(conv.getAssistantTable()) + rows, err := q.Select("tags").GroupBy("tags").Get() + if err != nil { + return nil, err + } + + tagSet := map[string]bool{} + for _, row := range rows { + if tags, ok := row["tags"].(string); ok && tags != "" { + var tagList []string + if err := jsoniter.UnmarshalFromString(tags, &tagList); err == nil { + for _, tag := range tagList { + tagSet[tag] = true + } + } + } + } + + // Convert map keys to slice + tags := make([]string, 0, len(tagSet)) + for tag := range tagSet { + tags = append(tags, tag) + } + return tags, nil +} diff --git a/neo/store/xun_test.go b/neo/store/xun_test.go index a41ae3fa..dc42de04 100644 --- a/neo/store/xun_test.go +++ b/neo/store/xun_test.go @@ -471,6 +471,9 @@ func TestXunAssistantCRUD(t *testing.T) { "avatar": "https://example.com/avatar.png", "connector": "openai", "description": "Test Description", + "path": "/assistants/test", + "sort": 100, + "built_in": true, "tags": tagsJSON, "options": optionsJSON, "mentionable": true, @@ -492,6 +495,9 @@ func TestXunAssistantCRUD(t *testing.T) { assert.Equal(t, "https://example.com/avatar.png", assistantData["avatar"]) assert.Equal(t, "openai", assistantData["connector"]) assert.Equal(t, "Test Description", assistantData["description"]) + assert.Equal(t, "/assistants/test", assistantData["path"]) + assert.Equal(t, int64(100), assistantData["sort"]) + assert.Equal(t, int64(1), assistantData["built_in"]) assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, assistantData["tags"]) assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, assistantData["options"]) assert.Equal(t, int64(1), assistantData["mentionable"]) @@ -504,6 +510,9 @@ func TestXunAssistantCRUD(t *testing.T) { "avatar": "https://example.com/avatar2.png", "connector": "openai", "description": "Test Description 2", + "path": "/assistants/test2", + "sort": 200, + "built_in": false, "tags": []string{"tag1", "tag2", "tag3"}, "options": map[string]interface{}{"model": "gpt-4"}, "prompts": []string{"prompt1", "prompt2"}, @@ -545,6 +554,9 @@ func TestXunAssistantCRUD(t *testing.T) { "type": "assistant", "connector": "openai", "description": "Test Description 3", + "path": nil, + "sort": 9999, + "built_in": false, "tags": nil, "options": nil, "prompts": nil, @@ -665,14 +677,232 @@ func TestXunAssistantCRUD(t *testing.T) { } } - // Test DeleteAssistant - err = store.DeleteAssistant(assistantID) + // Test non-existent assistant_id + resp, err = store.GetAssistants(AssistantFilter{ + AssistantID: "non-existent-id", + Page: 1, + PageSize: 10, + }) assert.Nil(t, err) - err = store.DeleteAssistant(assistant2ID) + assert.Equal(t, 0, len(resp.Data)) + + // Test filtering with select fields + resp, err = store.GetAssistants(AssistantFilter{ + Select: []string{"name", "description", "tags"}, + Page: 1, + PageSize: 10, + }) assert.Nil(t, err) - err = store.DeleteAssistant(assistant3ID) + // Verify only selected fields are returned + for _, item := range resp.Data { + // These fields should exist + assert.Contains(t, item, "name") + assert.Contains(t, item, "description") + assert.Contains(t, item, "tags") + // These fields should not exist + assert.NotContains(t, item, "options") + assert.NotContains(t, item, "prompts") + assert.NotContains(t, item, "flows") + assert.NotContains(t, item, "files") + assert.NotContains(t, item, "functions") + assert.NotContains(t, item, "permissions") + } + + // Test filtering with select fields and other filters combined + resp, err = store.GetAssistants(AssistantFilter{ + Tags: []string{"tag1"}, + Keywords: "Assistant", + Select: []string{"name", "tags"}, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + // Verify only selected fields are returned + for _, item := range resp.Data { + // These fields should exist + assert.Contains(t, item, "name") + assert.Contains(t, item, "tags") + // These fields should not exist + assert.NotContains(t, item, "description") + assert.NotContains(t, item, "options") + assert.NotContains(t, item, "prompts") + assert.NotContains(t, item, "flows") + assert.NotContains(t, item, "files") + assert.NotContains(t, item, "functions") + assert.NotContains(t, item, "permissions") + } + + // Test filtering with automated + automatedTrue := true + resp, err = store.GetAssistants(AssistantFilter{ + Automated: &automatedTrue, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Greater(t, len(resp.Data), 0) + + // Test filtering with mentionable + mentionableTrue := true + resp, err = store.GetAssistants(AssistantFilter{ + Mentionable: &mentionableTrue, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Greater(t, len(resp.Data), 0) + + // Test combined filters + resp, err = store.GetAssistants(AssistantFilter{ + Tags: []string{"tag1"}, + Keywords: "Assistant", + Connector: "openai", + Mentionable: &mentionableTrue, + Automated: &automatedTrue, + Page: 1, + PageSize: 10, + }) assert.Nil(t, err) + // Test filtering with built_in + builtInTrue := true + resp, err = store.GetAssistants(AssistantFilter{ + BuiltIn: &builtInTrue, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + for _, assistant := range resp.Data { + assert.Equal(t, int64(1), assistant["built_in"], "All assistants should be built-in") + } + + builtInFalse := false + resp, err = store.GetAssistants(AssistantFilter{ + BuiltIn: &builtInFalse, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + for _, assistant := range resp.Data { + assert.Equal(t, int64(0), assistant["built_in"], "All assistants should not be built-in") + } + + // Now test the delete operations + // First create some test data for delete operations + for i := 0; i < 5; i++ { + assistant := map[string]interface{}{ + "name": fmt.Sprintf("Delete Test Assistant %d", i), + "type": "assistant", + "connector": "openai", + "description": fmt.Sprintf("Delete Test Description %d", i), + "tags": []string{"delete-tag1", "delete-tag2"}, + "built_in": i%2 == 0, + "mentionable": true, + "automated": true, + } + _, err = store.SaveAssistant(assistant) + assert.Nil(t, err) + } + + // Test delete by connector + count, err := store.DeleteAssistants(AssistantFilter{ + Connector: "openai", + }) + assert.Nil(t, err) + assert.Greater(t, count, int64(0)) + + // Verify deletion + resp, err = store.GetAssistants(AssistantFilter{ + Connector: "openai", + }) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.Data)) + + // Create more test data for built_in test + for i := 0; i < 5; i++ { + assistant := map[string]interface{}{ + "name": fmt.Sprintf("Built-in Test Assistant %d", i), + "type": "assistant", + "connector": "openai", + "description": fmt.Sprintf("Built-in Test Description %d", i), + "tags": []string{"builtin-tag1", "builtin-tag2"}, + "built_in": true, + "mentionable": true, + "automated": true, + } + _, err = store.SaveAssistant(assistant) + assert.Nil(t, err) + } + + // Test delete by built_in status + builtInTrue = true + count, err = store.DeleteAssistants(AssistantFilter{ + BuiltIn: &builtInTrue, + }) + assert.Nil(t, err) + assert.Greater(t, count, int64(0)) + + // Verify deletion + resp, err = store.GetAssistants(AssistantFilter{ + BuiltIn: &builtInTrue, + }) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.Data)) + + // Create more test data for tags test + for i := 0; i < 5; i++ { + assistant := map[string]interface{}{ + "name": fmt.Sprintf("Tags Test Assistant %d", i), + "type": "assistant", + "connector": "openai", + "description": fmt.Sprintf("Tags Test Description %d", i), + "tags": []string{"tag1", "tag2"}, + "built_in": false, + "mentionable": true, + "automated": true, + } + _, err = store.SaveAssistant(assistant) + assert.Nil(t, err) + } + + // Test delete by tags + count, err = store.DeleteAssistants(AssistantFilter{ + Tags: []string{"tag1"}, + }) + assert.Nil(t, err) + assert.Greater(t, count, int64(0)) + + // Verify deletion + resp, err = store.GetAssistants(AssistantFilter{ + Tags: []string{"tag1"}, + }) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.Data)) + + // Create more test data for keywords test + for i := 0; i < 5; i++ { + assistant := map[string]interface{}{ + "name": fmt.Sprintf("Keywords Test Assistant %d", i), + "type": "assistant", + "connector": "openai", + "description": fmt.Sprintf("Keywords Test Description %d", i), + "tags": []string{"keyword-tag1", "keyword-tag2"}, + "built_in": false, + "mentionable": true, + "automated": true, + } + _, err = store.SaveAssistant(assistant) + assert.Nil(t, err) + } + + // Test delete by keywords + count, err = store.DeleteAssistants(AssistantFilter{ + Keywords: "Keywords Test", + }) + assert.Nil(t, err) + assert.Greater(t, count, int64(0)) + + // Verify all assistants are deleted resp, err = store.GetAssistants(AssistantFilter{}) assert.Nil(t, err) assert.Equal(t, 0, len(resp.Data)) @@ -725,6 +955,9 @@ func TestXunAssistantPagination(t *testing.T) { "connector": fmt.Sprintf("connector%d", i%3), "description": fmt.Sprintf("Description %d", i), "tags": tagsJSON, + "sort": 9999 - i, + "updated_at": time.Now().Add(time.Duration(-i) * time.Hour), + "built_in": i%2 == 0, "mentionable": mentionable, "automated": automated, } @@ -744,6 +977,25 @@ func TestXunAssistantPagination(t *testing.T) { assert.Equal(t, 2, resp.Next) assert.Equal(t, 0, resp.Prev) + // Verify sorting order (sort ASC, updated_at DESC) + for i := 1; i < len(resp.Data); i++ { + curr := resp.Data[i]["sort"].(int64) + prev := resp.Data[i-1]["sort"].(int64) + assert.True(t, curr >= prev, "Results should be sorted by sort ASC") + + // When sort values are equal, check updated_at if both values exist + if curr == prev { + currTime, currOk := resp.Data[i]["updated_at"].(time.Time) + prevTime, prevOk := resp.Data[i-1]["updated_at"].(time.Time) + + // Only compare times if both values exist + if currOk && prevOk { + assert.True(t, currTime.Before(prevTime) || currTime.Equal(prevTime), + "Results with same sort should be ordered by updated_at DESC") + } + } + } + // Test second page resp, err = store.GetAssistants(AssistantFilter{ Page: 2, @@ -811,7 +1063,30 @@ func TestXunAssistantPagination(t *testing.T) { assert.Nil(t, err) assert.Greater(t, len(resp.Data), 0) - // Test filtering by assistant_id + // Test filtering with built_in + builtInTrue := true + resp, err = store.GetAssistants(AssistantFilter{ + BuiltIn: &builtInTrue, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + for _, assistant := range resp.Data { + assert.Equal(t, int64(1), assistant["built_in"], "All assistants should be built-in") + } + + builtInFalse := false + resp, err = store.GetAssistants(AssistantFilter{ + BuiltIn: &builtInFalse, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + for _, assistant := range resp.Data { + assert.Equal(t, int64(0), assistant["built_in"], "All assistants should not be built-in") + } + + // Test assistant_id with other filters // First get an assistant_id from previous results firstAssistantID := resp.Data[0]["assistant_id"].(string) @@ -851,18 +1126,6 @@ func TestXunAssistantPagination(t *testing.T) { assert.Nil(t, err) assert.Equal(t, 0, len(resp.Data)) - // Test combined filters - resp, err = store.GetAssistants(AssistantFilter{ - Tags: []string{"tag0"}, - Keywords: "Assistant", - Connector: "connector0", - Mentionable: &mentionableTrue, - Automated: &automatedTrue, - Page: 1, - PageSize: 10, - }) - assert.Nil(t, err) - // Test filtering with select fields resp, err = store.GetAssistants(AssistantFilter{ Select: []string{"name", "description", "tags"}, @@ -909,4 +1172,152 @@ func TestXunAssistantPagination(t *testing.T) { assert.NotContains(t, item, "functions") assert.NotContains(t, item, "permissions") } + + // Test combined filters + resp, err = store.GetAssistants(AssistantFilter{ + Tags: []string{"tag0"}, + Keywords: "Assistant", + Connector: "connector0", + Mentionable: &mentionableTrue, + Automated: &automatedTrue, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + + // Now test the delete operations + // Test delete by connector + count, err := store.DeleteAssistants(AssistantFilter{ + Connector: "connector0", + }) + assert.Nil(t, err) + assert.Greater(t, count, int64(0)) + + // Verify deletion + resp, err = store.GetAssistants(AssistantFilter{ + Connector: "connector0", + }) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.Data)) + + // Test delete by built_in status + builtInTrue = true + count, err = store.DeleteAssistants(AssistantFilter{ + BuiltIn: &builtInTrue, + }) + assert.Nil(t, err) + assert.Greater(t, count, int64(0)) + + // Verify deletion + resp, err = store.GetAssistants(AssistantFilter{ + BuiltIn: &builtInTrue, + }) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.Data)) + + // Test delete by tags + count, err = store.DeleteAssistants(AssistantFilter{ + Tags: []string{"tag1"}, + }) + assert.Nil(t, err) + assert.Greater(t, count, int64(0)) + + // Verify deletion + resp, err = store.GetAssistants(AssistantFilter{ + Tags: []string{"tag1"}, + }) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.Data)) + + // Test delete by keywords + count, err = store.DeleteAssistants(AssistantFilter{ + Keywords: "Assistant", + }) + assert.Nil(t, err) + assert.Greater(t, count, int64(0)) + + // Verify all assistants are deleted + resp, err = store.GetAssistants(AssistantFilter{}) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.Data)) +} + +func TestGetAssistantTags(t *testing.T) { + + test.Prepare(t, config.Conf) + defer test.Clean() + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") + + store, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + if err != nil { + t.Fatal(err) + } + + // Create test assistants with tags + assistants := []map[string]interface{}{ + { + "assistant_id": "test-assistant-1", + "type": "assistant", + "connector": "test", + "tags": []string{"tag1", "tag2"}, + "name": "Test Assistant 1", + }, + { + "assistant_id": "test-assistant-2", + "type": "assistant", + "connector": "test", + "tags": []string{"tag2", "tag3"}, + "name": "Test Assistant 2", + }, + { + "assistant_id": "test-assistant-3", + "type": "assistant", + "connector": "test", + "tags": []string{"tag1", "tag3", "tag4"}, + "name": "Test Assistant 3", + }, + } + + // Save test assistants + for _, assistant := range assistants { + _, err := store.SaveAssistant(assistant) + if err != nil { + t.Fatal(err) + } + } + + // Get tags + tags, err := store.GetAssistantTags() + if err != nil { + t.Fatal(err) + } + + // Verify results + expectedTags := map[string]bool{ + "tag1": true, + "tag2": true, + "tag3": true, + "tag4": true, + } + + if len(tags) != len(expectedTags) { + t.Errorf("Expected %d tags, got %d", len(expectedTags), len(tags)) + } + + for _, tag := range tags { + if !expectedTags[tag] { + t.Errorf("Unexpected tag found: %s", tag) + } + } + + // Cleanup + for _, assistant := range assistants { + err := store.DeleteAssistant(assistant["assistant_id"].(string)) + if err != nil { + t.Fatal(err) + } + } }