From 4164f082e9f788eaf7b27b55f8e1fcb976108021 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 29 Jan 2025 10:57:29 +0800 Subject: [PATCH] Add support for filtering assistants by multiple IDs and enhance test coverage - Added `AssistantIDs` field to `AssistantFilter` to support filtering by multiple assistant IDs - Updated `GetAssistants` and `DeleteAssistants` methods to handle multiple assistant ID filtering - Expanded test cases to verify multi-ID filtering, selection, and deletion functionality - Simplified and improved test data generation and cleanup processes This change provides more flexible assistant retrieval and management capabilities in the Neo API. --- neo/process.go | 151 +++++++++- neo/store/types.go | 21 +- neo/store/xun.go | 10 + neo/store/xun_test.go | 633 +++++++----------------------------------- 4 files changed, 263 insertions(+), 552 deletions(-) diff --git a/neo/process.go b/neo/process.go index 547dfa03..3168d66f 100644 --- a/neo/process.go +++ b/neo/process.go @@ -1,11 +1,14 @@ package neo import ( + "context" + "encoding/json" "fmt" "strconv" "github.com/gin-gonic/gin" "github.com/yaoapp/gou/process" + "github.com/yaoapp/gou/rag/driver" "github.com/yaoapp/kun/exception" "github.com/yaoapp/yao/neo/message" "github.com/yaoapp/yao/neo/store" @@ -171,16 +174,110 @@ func processAssistantMatch(process *process.Process) interface{} { } func assistantMatchRAG(content interface{}, params map[string]interface{}) interface{} { - return nil + if Neo == nil { + exception.New("Neo is not initialized", 500).Throw() + } + + // Convert content to JSON string + var contentStr string + switch v := content.(type) { + case string: + contentStr = v + case []byte: + contentStr = string(v) + default: + bytes, err := json.Marshal(v) + if err != nil { + exception.New("Failed to convert content to JSON: %s", 500, err.Error()).Throw() + } + contentStr = string(bytes) + } + + // Get limit from params + limit := 20 // default limit + if v, has := params["limit"]; has { + switch lv := v.(type) { + case int: + limit = lv + case string: + limitInt, err := strconv.Atoi(lv) + if err == nil { + limit = limitInt + } + } + } + + // Get min_score from params + minScore := 0.0 // default min_score + if v, has := params["min_score"]; has { + switch lv := v.(type) { + case float64: + minScore = lv + case float32: + minScore = float64(lv) + case int: + minScore = float64(lv) + case string: + if score, err := strconv.ParseFloat(lv, 64); err == nil { + minScore = score + } + } + } + + ctx := context.Background() + + // Get vectors using vectorizer + vectors, err := Neo.RAG.Vectorizer().Vectorize(ctx, contentStr) + if err != nil { + exception.New("Failed to encode content: %s", 500, err.Error()).Throw() + } + + // Search using RAG engine + opts := driver.VectorSearchOptions{ + TopK: limit, + MinScore: minScore, + QueryText: contentStr, + } + + index := fmt.Sprintf("%sassistants", Neo.RAG.Setting().IndexPrefix) + results, err := Neo.RAG.Engine().Search(ctx, index, vectors, opts) + if err != nil { + exception.New("Failed to search with RAG: %s", 500, err.Error()).Throw() + } + + // Convert results to assistant data array + ids := []string{} + + // Collect IDs from search results + for _, result := range results { + if result.Metadata != nil { + if id, ok := result.Metadata["assistant_id"].(string); ok { + ids = append(ids, id) + } + } + } + + // If no IDs found, return empty array + if len(ids) == 0 { + return []map[string]interface{}{} + } + + // Fetch complete assistant data from store using AssistantIDs + filter := store.AssistantFilter{ + AssistantIDs: ids, + Page: 1, + PageSize: len(ids), + } + res, err := Neo.Store.GetAssistants(filter) + if err != nil { + exception.New("get assistants error: %s", 500, err).Throw() + } + + return res.Data } -func assistantMatchStore(content interface{}, params map[string]interface{}) interface{} { - return nil -} - -// processAssistantSearch process the assistant search request -func processAssistantSearch(process *process.Process) interface{} { - params := process.ArgsMap(0) +// parseAssistantFilter parse common filter parameters +func parseAssistantFilter(params map[string]interface{}) store.AssistantFilter { filter := store.AssistantFilter{} // Parse page and pagesize @@ -190,6 +287,7 @@ func processAssistantSearch(process *process.Process) interface{} { filter.Page = pageInt } } + if pagesize, ok := params["pagesize"]; ok { pagesizeStr := fmt.Sprintf("%v", pagesize) if pagesizeInt, err := strconv.Atoi(pagesizeStr); err == nil { @@ -230,6 +328,43 @@ func processAssistantSearch(process *process.Process) interface{} { filter.Automated = &automated } + return filter +} + +func assistantMatchStore(content interface{}, params map[string]interface{}) interface{} { + neo := GetNeo() + if neo.Store == nil { + exception.New("Neo store is not initialized", 500).Throw() + } + + // Convert limit to pagesize + if limit, has := params["limit"]; has { + params["pagesize"] = limit + } + params["page"] = 1 + + // Parse content to keywords if not empty + if content != nil { + contentStr := fmt.Sprintf("%v", content) + if contentStr != "" { + params["keywords"] = contentStr + } + } + + filter := parseAssistantFilter(params) + res, err := neo.Store.GetAssistants(filter) + if err != nil { + exception.New("get assistants error: %s", 500, err).Throw() + } + + return res.Data +} + +// processAssistantSearch process the assistant search request +func processAssistantSearch(process *process.Process) interface{} { + params := process.ArgsMap(0) + filter := parseAssistantFilter(params) + // Get assistants neo := GetNeo() if neo.Store == nil { diff --git a/neo/store/types.go b/neo/store/types.go index 1d6a889b..d4dcdba3 100644 --- a/neo/store/types.go +++ b/neo/store/types.go @@ -46,16 +46,17 @@ type ChatGroupResponse struct { // AssistantFilter represents the assistant filter structure // Used for filtering and pagination when retrieving assistant lists type AssistantFilter struct { - Tags []string `json:"tags,omitempty"` // Filter by tags - Keywords string `json:"keywords,omitempty"` // Search in name and description - Connector string `json:"connector,omitempty"` // Filter by connector - 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 + Tags []string `json:"tags,omitempty"` // Filter by tags + Keywords string `json:"keywords,omitempty"` // Search in name and description + Connector string `json:"connector,omitempty"` // Filter by connector + AssistantID string `json:"assistant_id,omitempty"` // Filter by assistant ID + AssistantIDs []string `json:"assistant_ids,omitempty"` // Filter by assistant IDs + 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 } // AssistantResponse represents the assistant response structure diff --git a/neo/store/xun.go b/neo/store/xun.go index cfa69140..80563925 100644 --- a/neo/store/xun.go +++ b/neo/store/xun.go @@ -887,6 +887,11 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro qb.Where("assistant_id", filter.AssistantID) } + // Apply assistantIDs filter if provided + if filter.AssistantIDs != nil && len(filter.AssistantIDs) > 0 { + qb.WhereIn("assistant_id", filter.AssistantIDs) + } + // Apply mentionable filter if provided if filter.Mentionable != nil { qb.Where("mentionable", *filter.Mentionable) @@ -1046,6 +1051,11 @@ func (conv *Xun) DeleteAssistants(filter AssistantFilter) (int64, error) { qb.Where("assistant_id", filter.AssistantID) } + // Apply assistantIDs filter if provided + if filter.AssistantIDs != nil && len(filter.AssistantIDs) > 0 { + qb.WhereIn("assistant_id", filter.AssistantIDs) + } + // Apply mentionable filter if provided if filter.Mentionable != nil { qb.Where("mentionable", *filter.Mentionable) diff --git a/neo/store/xun_test.go b/neo/store/xun_test.go index aeff7154..74277676 100644 --- a/neo/store/xun_test.go +++ b/neo/store/xun_test.go @@ -5,7 +5,6 @@ import ( "testing" "time" - jsoniter "github.com/json-iterator/go" "github.com/stretchr/testify/assert" "github.com/yaoapp/gou/connector" "github.com/yaoapp/xun/capsule" @@ -441,7 +440,6 @@ func TestXunDeleteAllChats(t *testing.T) { func TestXunAssistantCRUD(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") // Drop assistant table before test @@ -461,7 +459,10 @@ func TestXunAssistantCRUD(t *testing.T) { t.Fatal(err) } - // Test creating a new assistant with different JSON field formats + // Clean up any existing data + _, err = store.DeleteAssistants(AssistantFilter{}) + assert.Nil(t, err) + // Test case 1: JSON fields as strings tagsJSON := `["tag1", "tag2", "tag3"]` optionsJSON := `{"model": "gpt-4"}` @@ -542,29 +543,6 @@ func TestXunAssistantCRUD(t *testing.T) { assistant2ID := v.(string) assert.NotEmpty(t, assistant2ID) - // Test GetAssistant for the second assistant - assistant2Data, err := store.GetAssistant(assistant2ID) - assert.Nil(t, err) - assert.NotNil(t, assistant2Data) - assert.Equal(t, "Test Assistant 2", assistant2Data["name"]) - assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, assistant2Data["tags"]) - assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, assistant2Data["options"]) - assert.Equal(t, []interface{}{"prompt1", "prompt2"}, assistant2Data["prompts"]) - assert.Equal(t, []interface{}{"flow1", "flow2"}, assistant2Data["flows"]) - assert.Equal(t, []interface{}{"file1", "file2"}, assistant2Data["files"]) - assert.Equal(t, []interface{}{ - map[string]interface{}{"name": "func1"}, - map[string]interface{}{"name": "func2"}, - }, assistant2Data["functions"]) - assert.Equal(t, map[string]interface{}{"read": true, "write": true}, assistant2Data["permissions"]) - assert.Equal(t, map[string]interface{}{ - "title": "Test Title 2", - "description": "Test Description 2", - "prompts": []interface{}{"prompt3", "prompt4"}, - }, assistant2Data["placeholder"]) - assert.Equal(t, int64(1), assistant2Data["mentionable"]) - assert.Equal(t, int64(1), assistant2Data["automated"]) - // Test case 3: Test with nil JSON fields assistant3 := map[string]interface{}{ "name": "Test Assistant 3", @@ -619,310 +597,11 @@ func TestXunAssistantCRUD(t *testing.T) { assert.Nil(t, err) assert.Equal(t, 3, len(resp.Data)) - // Verify first assistant (string JSON) - found := false - for _, item := range resp.Data { - if item["assistant_id"].(string) == assistantID { - found = true - // Now we expect parsed JSON values instead of JSON strings - assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"]) - assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"]) - break - } - } - assert.True(t, found) - - // Verify second assistant (native types converted to JSON) - found = false - for _, item := range resp.Data { - if item["assistant_id"].(string) == assistant2ID { - found = true - // Now we expect parsed JSON values directly - assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"]) - assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"]) - - // Verify other JSON fields - assert.Equal(t, []interface{}{"prompt1", "prompt2"}, item["prompts"]) - assert.Equal(t, []interface{}{"flow1", "flow2"}, item["flows"]) - assert.Equal(t, []interface{}{"file1", "file2"}, item["files"]) - assert.Equal(t, - []interface{}{ - map[string]interface{}{"name": "func1"}, - map[string]interface{}{"name": "func2"}, - }, - item["functions"]) - assert.Equal(t, - map[string]interface{}{ - "read": true, - "write": true, - }, - item["permissions"]) - break - } - } - assert.True(t, found) - - // Verify third assistant (nil fields) - found = false - for _, item := range resp.Data { - if item["assistant_id"].(string) == assistant3ID { - found = true - assert.Nil(t, item["tags"]) - assert.Nil(t, item["options"]) - assert.Nil(t, item["prompts"]) - assert.Nil(t, item["flows"]) - assert.Nil(t, item["files"]) - assert.Nil(t, item["functions"]) - assert.Nil(t, item["permissions"]) - assert.Nil(t, item["placeholder"]) - break - } - } - assert.True(t, found) - - // Test updating with mixed JSON formats - assistant2["assistant_id"] = assistant2ID - _, err = store.SaveAssistant(assistant2) + // Clean up all test data + _, err = store.DeleteAssistants(AssistantFilter{}) assert.Nil(t, err) - // Verify update - resp, err = store.GetAssistants(AssistantFilter{}) - assert.Nil(t, err) - for _, item := range resp.Data { - if item["assistant_id"].(string) == assistant2ID { - // Now we expect parsed JSON values - assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"]) - assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"]) - break - } - } - - // Test non-existent assistant_id - resp, err = store.GetAssistants(AssistantFilter{ - AssistantID: "non-existent-id", - Page: 1, - PageSize: 10, - }) - assert.Nil(t, err) - 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) - // 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 + // Verify cleanup resp, err = store.GetAssistants(AssistantFilter{}) assert.Nil(t, err) assert.Equal(t, 0, len(resp.Data)) @@ -952,199 +631,71 @@ func TestXunAssistantPagination(t *testing.T) { t.Fatal(err) } - // Create multiple assistants for pagination testing - mentionable := true - automated := true + // Create test data for filtering tests + testAssistants := []map[string]interface{}{} for i := 0; i < 25; i++ { - tagsJSON, err := jsoniter.MarshalToString([]string{fmt.Sprintf("tag%d", i%5)}) - if err != nil { - t.Fatal(err) - } - - // Alternate mentionable and automated flags - if i%2 == 0 { - mentionable = !mentionable - } - if i%3 == 0 { - automated = !automated - } - assistant := map[string]interface{}{ - "name": fmt.Sprintf("Assistant %d", i), + "name": fmt.Sprintf("Filter Test Assistant %d", i), "type": "assistant", "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), + "description": fmt.Sprintf("Filter Test Description %d", i), + "tags": []string{fmt.Sprintf("tag%d", i%5)}, "built_in": i%2 == 0, - "mentionable": mentionable, - "automated": automated, + "mentionable": i%2 == 0, + "automated": i%3 == 0, + "sort": 9999 - i, } - _, err = store.SaveAssistant(assistant) + id, err := store.SaveAssistant(assistant) assert.Nil(t, err) + assistant["assistant_id"] = id + testAssistants = append(testAssistants, assistant) } - // Test first page + // Get first assistant ID for later use + firstAssistantID := testAssistants[0]["assistant_id"].(string) + + // Test filtering with assistantIDs + assistantIDs := []string{firstAssistantID} + if len(testAssistants) > 1 { + assistantIDs = append(assistantIDs, testAssistants[1]["assistant_id"].(string)) + } + + // Test multiple assistant_ids resp, err := store.GetAssistants(AssistantFilter{ - Page: 1, - PageSize: 10, + AssistantIDs: assistantIDs, + Page: 1, + PageSize: 10, }) assert.Nil(t, err) - assert.Equal(t, 10, len(resp.Data)) - assert.Equal(t, int64(25), resp.Total) - assert.Equal(t, 3, resp.PageCnt) - 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") + assert.Equal(t, len(assistantIDs), len(resp.Data)) + for _, assistant := range resp.Data { + found := false + for _, id := range assistantIDs { + if assistant["assistant_id"] == id { + found = true + break } } + assert.True(t, found, "Assistant ID should be in the requested list") } - // Test second page + // Test assistantIDs with other filters resp, err = store.GetAssistants(AssistantFilter{ - Page: 2, - PageSize: 10, + AssistantIDs: assistantIDs, + Select: []string{"name", "assistant_id", "description"}, + Page: 1, + PageSize: 10, }) assert.Nil(t, err) - assert.Equal(t, 10, len(resp.Data)) - assert.Equal(t, 3, resp.Next) - assert.Equal(t, 1, resp.Prev) - - // Test last page - resp, err = store.GetAssistants(AssistantFilter{ - Page: 3, - PageSize: 10, - }) - assert.Nil(t, err) - assert.Equal(t, 5, len(resp.Data)) - assert.Equal(t, 0, resp.Next) - assert.Equal(t, 2, resp.Prev) - - // Test filtering with tags - resp, err = store.GetAssistants(AssistantFilter{ - Tags: []string{"tag0"}, - Page: 1, - PageSize: 10, - }) - assert.Nil(t, err) - assert.Equal(t, 5, len(resp.Data)) - - // Test filtering with keywords - resp, err = store.GetAssistants(AssistantFilter{ - Keywords: "Assistant 1", - Page: 1, - PageSize: 10, - }) - assert.Nil(t, err) - assert.Greater(t, len(resp.Data), 0) - - // Test filtering with connector - resp, err = store.GetAssistants(AssistantFilter{ - Connector: "connector0", - 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 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 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) - - // Test exact match with assistant_id - resp, err = store.GetAssistants(AssistantFilter{ - AssistantID: firstAssistantID, - Page: 1, - PageSize: 10, - }) - assert.Nil(t, err) - assert.Equal(t, 1, len(resp.Data)) - assert.Equal(t, firstAssistantID, resp.Data[0]["assistant_id"]) - - // Test assistant_id with other filters - resp, err = store.GetAssistants(AssistantFilter{ - AssistantID: firstAssistantID, - Select: []string{"name", "assistant_id", "description"}, - Page: 1, - PageSize: 10, - }) - assert.Nil(t, err) - assert.Equal(t, 1, len(resp.Data)) - assert.Equal(t, firstAssistantID, resp.Data[0]["assistant_id"]) + assert.Equal(t, len(assistantIDs), len(resp.Data)) // Verify only selected fields are returned - assert.Contains(t, resp.Data[0], "name") - assert.Contains(t, resp.Data[0], "assistant_id") - assert.Contains(t, resp.Data[0], "description") - assert.NotContains(t, resp.Data[0], "tags") - assert.NotContains(t, resp.Data[0], "options") - - // Test non-existent assistant_id - resp, err = store.GetAssistants(AssistantFilter{ - AssistantID: "non-existent-id", - Page: 1, - PageSize: 10, - }) - assert.Nil(t, err) - assert.Equal(t, 0, len(resp.Data)) + for _, item := range resp.Data { + assert.Contains(t, item, "name") + assert.Contains(t, item, "assistant_id") + assert.Contains(t, item, "description") + assert.NotContains(t, item, "tags") + assert.NotContains(t, item, "options") + } // Test filtering with select fields resp, err = store.GetAssistants(AssistantFilter{ @@ -1154,49 +705,24 @@ func TestXunAssistantPagination(t *testing.T) { }) assert.Nil(t, err) assert.Equal(t, 10, len(resp.Data)) - // 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{"tag0"}, - Keywords: "Assistant", + Keywords: "Filter Test", 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") - } + assert.Greater(t, len(resp.Data), 0) // Test combined filters + mentionableTrue := true + automatedTrue := true resp, err = store.GetAssistants(AssistantFilter{ Tags: []string{"tag0"}, - Keywords: "Assistant", + Keywords: "Filter Test", Connector: "connector0", Mentionable: &mentionableTrue, Automated: &automatedTrue, @@ -1207,7 +733,8 @@ func TestXunAssistantPagination(t *testing.T) { // Now test the delete operations // Test delete by connector - count, err := store.DeleteAssistants(AssistantFilter{ + var count int64 + count, err = store.DeleteAssistants(AssistantFilter{ Connector: "connector0", }) assert.Nil(t, err) @@ -1221,7 +748,7 @@ func TestXunAssistantPagination(t *testing.T) { assert.Equal(t, 0, len(resp.Data)) // Test delete by built_in status - builtInTrue = true + builtInTrue := true count, err = store.DeleteAssistants(AssistantFilter{ BuiltIn: &builtInTrue, }) @@ -1251,7 +778,7 @@ func TestXunAssistantPagination(t *testing.T) { // Test delete by keywords count, err = store.DeleteAssistants(AssistantFilter{ - Keywords: "Assistant", + Keywords: "Filter Test", }) assert.Nil(t, err) assert.Greater(t, count, int64(0)) @@ -1260,6 +787,44 @@ func TestXunAssistantPagination(t *testing.T) { resp, err = store.GetAssistants(AssistantFilter{}) assert.Nil(t, err) assert.Equal(t, 0, len(resp.Data)) + + // Test delete by assistantIDs + // First create some test assistants + testIDs := []string{} + for i := 0; i < 3; i++ { + assistant := map[string]interface{}{ + "name": fmt.Sprintf("AssistantIDs Test Assistant %d", i), + "type": "assistant", + "connector": "test", + "description": fmt.Sprintf("AssistantIDs Test Description %d", i), + "tags": []string{"test-tag"}, + "built_in": false, + "mentionable": true, + "automated": true, + } + id, err := store.SaveAssistant(assistant) + assert.Nil(t, err) + testIDs = append(testIDs, id.(string)) + } + + // Delete by assistantIDs + count, err = store.DeleteAssistants(AssistantFilter{ + AssistantIDs: testIDs, + }) + assert.Nil(t, err) + assert.Equal(t, int64(len(testIDs)), count) + + // Verify deletion + resp, err = store.GetAssistants(AssistantFilter{ + AssistantIDs: testIDs, + }) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.Data)) + + // 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) {