From af65a2d5fa6545014e3ef65152665c64c74c67a8 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 29 Dec 2024 12:57:17 +0800 Subject: [PATCH] Refactor assistant management and enhance API response structure - Updated the assistant management functions to utilize a unified response structure, replacing the previous 'Items' field with 'Data' for consistency across the API. - Refactored the handling of assistant mentions and details to streamline data processing, improving code readability and maintainability. - Enhanced the AssistantResponse type to include pagination fields directly, facilitating better management of assistant data retrieval. - Implemented new methods for adding, saving, deleting, and searching assistants, ensuring robust functionality across different storage backends. - Updated tests to reflect changes in the assistant management logic and response structure, ensuring comprehensive coverage and reliability. --- neo/api.go | 28 ++- neo/conversation/mongo.go | 19 +- neo/conversation/redis.go | 19 +- neo/conversation/types.go | 12 +- neo/conversation/weaviate.go | 20 +- neo/conversation/xun.go | 102 +++++++-- neo/conversation/xun_test.go | 64 +++--- neo/process.go | 138 ++++++++++++- neo/process_test.go | 390 +++++++++++++++++++++++++++++++++++ test/utils.go | 45 ++++ 10 files changed, 737 insertions(+), 100 deletions(-) create mode 100644 neo/process_test.go diff --git a/neo/api.go b/neo/api.go index 0db39ee8..9da02e91 100644 --- a/neo/api.go +++ b/neo/api.go @@ -491,16 +491,14 @@ func (neo *DSL) handleMentions(c *gin.Context) { // Convert assistants to mentions mentions := []Mention{} - for _, item := range response.Items { - if assistant, ok := item.(map[string]interface{}); ok { - mention := Mention{ - ID: assistant["assistant_id"].(string), - Name: assistant["name"].(string), - Type: assistant["type"].(string), - Avatar: assistant["avatar"].(string), - } - mentions = append(mentions, mention) + for _, item := range response.Data { + mention := Mention{ + ID: item["assistant_id"].(string), + Name: item["name"].(string), + Type: item["type"].(string), + Avatar: item["avatar"].(string), } + mentions = append(mentions, mention) } c.JSON(200, map[string]interface{}{"data": mentions}) @@ -872,7 +870,7 @@ func (neo *DSL) handleAssistantList(c *gin.Context) { return } - c.JSON(200, map[string]interface{}{"data": response}) + c.JSON(200, response) c.Done() } @@ -899,12 +897,10 @@ func (neo *DSL) handleAssistantDetail(c *gin.Context) { // Find the assistant by ID var assistant map[string]interface{} - for _, item := range response.Items { - if a, ok := item.(map[string]interface{}); ok { - if id, ok := a["id"].(string); ok && id == assistantID { - assistant = a - break - } + for _, item := range response.Data { + if id, ok := item["id"].(string); ok && id == assistantID { + assistant = item + break } } diff --git a/neo/conversation/mongo.go b/neo/conversation/mongo.go index b406ffac..c96fddf6 100644 --- a/neo/conversation/mongo.go +++ b/neo/conversation/mongo.go @@ -1,7 +1,5 @@ package conversation -import "github.com/yaoapp/xun" - // Mongo conversation type Mongo struct{} @@ -74,15 +72,12 @@ func (conv *Mongo) DeleteAssistant(assistantID string) error { // GetAssistants retrieves assistants with pagination and tag filtering func (conv *Mongo) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { return &AssistantResponse{ - P: xun.P{ - Items: []interface{}{}, - Total: 0, - TotalPages: 0, - PageSize: filter.PageSize, - CurrentPage: filter.Page, - NextPage: 0, - PreviousPage: 0, - LastPage: 0, - }, + Data: []map[string]interface{}{}, + Page: filter.Page, + PageSize: filter.PageSize, + PageCnt: 0, + Next: 0, + Prev: 0, + Total: 0, }, nil } diff --git a/neo/conversation/redis.go b/neo/conversation/redis.go index 9aef5998..7f88dbf5 100644 --- a/neo/conversation/redis.go +++ b/neo/conversation/redis.go @@ -1,7 +1,5 @@ package conversation -import "github.com/yaoapp/xun" - // Redis conversation type Redis struct{} @@ -74,15 +72,12 @@ func (conv *Redis) DeleteAssistant(assistantID string) error { // GetAssistants retrieves assistants with pagination and tag filtering func (conv *Redis) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { return &AssistantResponse{ - P: xun.P{ - Items: []interface{}{}, - Total: 0, - TotalPages: 0, - PageSize: filter.PageSize, - CurrentPage: filter.Page, - NextPage: 0, - PreviousPage: 0, - LastPage: 0, - }, + Data: []map[string]interface{}{}, + Page: filter.Page, + PageSize: filter.PageSize, + PageCnt: 0, + Next: 0, + Prev: 0, + Total: 0, }, nil } diff --git a/neo/conversation/types.go b/neo/conversation/types.go index b85c8eeb..87c0176a 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -1,7 +1,5 @@ package conversation -import "github.com/yaoapp/xun" - // Setting represents the conversation configuration structure // Used to configure basic conversation parameters including connector, user field, table name, etc. type Setting struct { @@ -58,9 +56,15 @@ type AssistantFilter struct { } // AssistantResponse represents the assistant response structure -// Inherits from xun.P, used for returning paginated assistant lists +// Used for returning paginated assistant lists type AssistantResponse struct { - xun.P + Data []map[string]interface{} `json:"data"` // The paginated data + Page int `json:"page"` // Current page number + PageSize int `json:"pagesize"` // Number of items per page + PageCnt int `json:"pagecnt"` // Total number of pages + Next int `json:"next"` // Next page number + Prev int `json:"prev"` // Previous page number + Total int64 `json:"total"` // Total number of items } // Conversation defines the conversation storage interface diff --git a/neo/conversation/weaviate.go b/neo/conversation/weaviate.go index 932aa22a..0b572111 100644 --- a/neo/conversation/weaviate.go +++ b/neo/conversation/weaviate.go @@ -1,7 +1,5 @@ package conversation -import "github.com/yaoapp/xun" - // Weaviate Database conversation type Weaviate struct{} @@ -74,15 +72,13 @@ func (conv *Weaviate) DeleteAssistant(assistantID string) error { // GetAssistants retrieves assistants with pagination and tag filtering func (conv *Weaviate) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { return &AssistantResponse{ - P: xun.P{ - Items: []interface{}{}, - Total: 0, - TotalPages: 0, - PageSize: filter.PageSize, - CurrentPage: filter.Page, - NextPage: 0, - PreviousPage: 0, - LastPage: 0, - }, + + Data: []map[string]interface{}{}, + Page: filter.Page, + PageSize: filter.PageSize, + PageCnt: 0, + Next: 0, + Prev: 0, + Total: 0, }, nil } diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 5e439e73..752cda0e 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -662,35 +662,47 @@ func (conv *Xun) DeleteAllChats(sid string) error { return err } -// SaveAssistant creates or updates an assistant +// SaveAssistant saves assistant information func (conv *Xun) SaveAssistant(assistant map[string]interface{}) error { - assistantID, ok := assistant["assistant_id"].(string) - if !ok || assistantID == "" { - assistantID = uuid.New().String() - assistant["assistant_id"] = assistantID + // Validate required fields + requiredFields := []string{"name", "type", "connector"} + for _, field := range requiredFields { + if _, ok := assistant[field]; !ok { + return fmt.Errorf("field %s is required", field) + } + if assistant[field] == nil || assistant[field] == "" { + return fmt.Errorf("field %s cannot be empty", field) + } + } + + // Validate tags format + if tags, ok := assistant["tags"].(string); ok { + log.Trace("Saving assistant with tags: %s", tags) + } + + // Generate assistant_id if not provided + if _, ok := assistant["assistant_id"]; !ok { + assistant["assistant_id"] = uuid.New().String() } // Check if assistant exists exists, err := conv.query.New(). Table(conv.getAssistantTable()). - Where("assistant_id", assistantID). + Where("assistant_id", assistant["assistant_id"]). Exists() if err != nil { return err } - now := time.Now() - assistant["updated_at"] = now - + // Update or insert if exists { - // Update existing assistant + assistant["updated_at"] = time.Now() _, err = conv.query.New(). Table(conv.getAssistantTable()). - Where("assistant_id", assistantID). + Where("assistant_id", assistant["assistant_id"]). Update(assistant) } else { - // Create new assistant - assistant["created_at"] = now + assistant["created_at"] = time.Now() err = conv.query.New(). Table(conv.getAssistantTable()). Insert(assistant) @@ -701,7 +713,20 @@ func (conv *Xun) SaveAssistant(assistant map[string]interface{}) error { // DeleteAssistant deletes an assistant by assistant_id func (conv *Xun) DeleteAssistant(assistantID string) error { - _, err := conv.query.New(). + // Check if assistant exists + exists, err := conv.query.New(). + Table(conv.getAssistantTable()). + Where("assistant_id", assistantID). + Exists() + if err != nil { + return err + } + + if !exists { + return fmt.Errorf("assistant %s not found", assistantID) + } + + _, err = conv.query.New(). Table(conv.getAssistantTable()). Where("assistant_id", assistantID). Delete() @@ -717,10 +742,13 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro if filter.Tags != nil && len(filter.Tags) > 0 { qb.Where(func(qb query.Query) { for i, tag := range filter.Tags { + // For each tag, we need to match it as part of a JSON array + // This will match both single tag arrays ["tag1"] and multi-tag arrays ["tag1","tag2"] + pattern := fmt.Sprintf("%%\"%s\"%%", tag) if i == 0 { - qb.Where("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag)) + qb.Where("tags", "like", pattern) } else { - qb.OrWhere("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag)) + qb.OrWhere("tags", "like", pattern) } } }) @@ -757,12 +785,46 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro filter.Page = 1 } - // Get paginated results - paginator, err := qb.OrderBy("created_at", "desc"). - Paginate(filter.PageSize, filter.Page) + // Get total count + total, err := qb.Clone().Count() if err != nil { return nil, err } - return &AssistantResponse{P: paginator}, nil + // Calculate pagination + offset := (filter.Page - 1) * filter.PageSize + totalPages := int(math.Ceil(float64(total) / float64(filter.PageSize))) + nextPage := filter.Page + 1 + if nextPage > totalPages { + nextPage = 0 + } + prevPage := filter.Page - 1 + if prevPage < 1 { + prevPage = 0 + } + + // Get paginated results + rows, err := qb.OrderBy("created_at", "desc"). + Offset(offset). + Limit(filter.PageSize). + Get() + if err != nil { + return nil, err + } + + // Convert rows to map slice + data := make([]map[string]interface{}, len(rows)) + for i, row := range rows { + data[i] = row + } + + return &AssistantResponse{ + Data: data, + Page: filter.Page, + PageSize: filter.PageSize, + PageCnt: totalPages, + Next: nextPage, + Prev: prevPage, + Total: total, + }, nil } diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go index 90eacf74..a08c800b 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -8,7 +8,6 @@ import ( jsoniter "github.com/json-iterator/go" "github.com/stretchr/testify/assert" "github.com/yaoapp/gou/connector" - "github.com/yaoapp/xun" "github.com/yaoapp/xun/capsule" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/test" @@ -36,6 +35,9 @@ func TestNewXunDefault(t *testing.T) { t.Fatal(err) } + // Add a small delay to ensure table is created + time.Sleep(100 * time.Millisecond) + conv, err := NewXun(Setting{ Connector: "default", Table: "__unit_test_conversation", @@ -123,6 +125,9 @@ func TestNewXunConnector(t *testing.T) { sch.DropTableIfExists("__unit_test_conversation_chat") sch.DropTableIfExists("__unit_test_conversation_assistant") + // Add a small delay to ensure table is created + time.Sleep(100 * time.Millisecond) + conv, err := NewXun(Setting{ Connector: "mysql", Table: "__unit_test_conversation", @@ -462,6 +467,9 @@ func TestXunAssistantCRUD(t *testing.T) { t.Fatal(err) } + // Add a small delay to ensure table is created + time.Sleep(100 * time.Millisecond) + conv, err := NewXun(Setting{ Connector: "default", Table: "__unit_test_conversation", @@ -507,56 +515,56 @@ func TestXunAssistantCRUD(t *testing.T) { // Test GetAssistants with no filter resp, err := conv.GetAssistants(AssistantFilter{}) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.P.Items)) + assert.Equal(t, 1, len(resp.Data)) // Test GetAssistants with tag filter (single tag) resp, err = conv.GetAssistants(AssistantFilter{ Tags: []string{"tag1"}, }) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.P.Items)) + assert.Equal(t, 1, len(resp.Data)) // Test GetAssistants with tag filter (multiple tags) resp, err = conv.GetAssistants(AssistantFilter{ Tags: []string{"tag1", "tag4"}, }) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.P.Items)) + assert.Equal(t, 1, len(resp.Data)) // Test GetAssistants with non-existent tag resp, err = conv.GetAssistants(AssistantFilter{ Tags: []string{"nonexistent"}, }) assert.Nil(t, err) - assert.Equal(t, 0, len(resp.P.Items)) + assert.Equal(t, 0, len(resp.Data)) // Test GetAssistants with keyword filter resp, err = conv.GetAssistants(AssistantFilter{ Keywords: "Test", }) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.P.Items)) + assert.Equal(t, 1, len(resp.Data)) // Test GetAssistants with connector filter resp, err = conv.GetAssistants(AssistantFilter{ Connector: "openai", }) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.P.Items)) + assert.Equal(t, 1, len(resp.Data)) // Test GetAssistants with mentionable filter resp, err = conv.GetAssistants(AssistantFilter{ Mentionable: &mentionable, }) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.P.Items)) + assert.Equal(t, 1, len(resp.Data)) // Test GetAssistants with automated filter resp, err = conv.GetAssistants(AssistantFilter{ Automated: &automated, }) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.P.Items)) + assert.Equal(t, 1, len(resp.Data)) // Test GetAssistants with combined filters resp, err = conv.GetAssistants(AssistantFilter{ @@ -567,7 +575,7 @@ func TestXunAssistantCRUD(t *testing.T) { Tags: []string{"tag1"}, }) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.P.Items)) + assert.Equal(t, 1, len(resp.Data)) // Test SaveAssistant (Update) assistant["name"] = "Updated Assistant" @@ -576,8 +584,8 @@ func TestXunAssistantCRUD(t *testing.T) { resp, err = conv.GetAssistants(AssistantFilter{}) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.P.Items)) - item := resp.P.Items[0].(xun.R) + assert.Equal(t, 1, len(resp.Data)) + item := resp.Data[0] assert.Equal(t, "Updated Assistant", item["name"]) // Test DeleteAssistant @@ -586,13 +594,14 @@ func TestXunAssistantCRUD(t *testing.T) { resp, err = conv.GetAssistants(AssistantFilter{}) assert.Nil(t, err) - assert.Equal(t, 0, len(resp.P.Items)) + assert.Equal(t, 0, len(resp.Data)) } func TestXunAssistantPagination(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_chat") defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") // Drop assistant table before test @@ -601,6 +610,9 @@ func TestXunAssistantPagination(t *testing.T) { t.Fatal(err) } + // Add a small delay to ensure table is created + time.Sleep(100 * time.Millisecond) + conv, err := NewXun(Setting{ Connector: "default", Table: "__unit_test_conversation", @@ -645,9 +657,11 @@ func TestXunAssistantPagination(t *testing.T) { PageSize: 10, }) assert.Nil(t, err) - assert.Equal(t, 10, len(resp.P.Items)) - assert.Equal(t, 25, resp.P.Total) - assert.Equal(t, 3, resp.P.LastPage) + 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) // Test second page resp, err = conv.GetAssistants(AssistantFilter{ @@ -655,7 +669,9 @@ func TestXunAssistantPagination(t *testing.T) { PageSize: 10, }) assert.Nil(t, err) - assert.Equal(t, 10, len(resp.P.Items)) + assert.Equal(t, 10, len(resp.Data)) + assert.Equal(t, 3, resp.Next) + assert.Equal(t, 1, resp.Prev) // Test last page resp, err = conv.GetAssistants(AssistantFilter{ @@ -663,7 +679,9 @@ func TestXunAssistantPagination(t *testing.T) { PageSize: 10, }) assert.Nil(t, err) - assert.Equal(t, 5, len(resp.P.Items)) + 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 = conv.GetAssistants(AssistantFilter{ @@ -672,7 +690,7 @@ func TestXunAssistantPagination(t *testing.T) { PageSize: 10, }) assert.Nil(t, err) - assert.Equal(t, 5, len(resp.P.Items)) + assert.Equal(t, 5, len(resp.Data)) // Test filtering with keywords resp, err = conv.GetAssistants(AssistantFilter{ @@ -681,7 +699,7 @@ func TestXunAssistantPagination(t *testing.T) { PageSize: 10, }) assert.Nil(t, err) - assert.Greater(t, len(resp.P.Items), 0) + assert.Greater(t, len(resp.Data), 0) // Test filtering with connector resp, err = conv.GetAssistants(AssistantFilter{ @@ -690,7 +708,7 @@ func TestXunAssistantPagination(t *testing.T) { PageSize: 10, }) assert.Nil(t, err) - assert.Greater(t, len(resp.P.Items), 0) + assert.Greater(t, len(resp.Data), 0) // Test filtering with mentionable mentionableTrue := true @@ -700,7 +718,7 @@ func TestXunAssistantPagination(t *testing.T) { PageSize: 10, }) assert.Nil(t, err) - assert.Greater(t, len(resp.P.Items), 0) + assert.Greater(t, len(resp.Data), 0) // Test filtering with automated automatedTrue := true @@ -710,7 +728,7 @@ func TestXunAssistantPagination(t *testing.T) { PageSize: 10, }) assert.Nil(t, err) - assert.Greater(t, len(resp.P.Items), 0) + assert.Greater(t, len(resp.Data), 0) // Test combined filters resp, err = conv.GetAssistants(AssistantFilter{ diff --git a/neo/process.go b/neo/process.go index 598b40ad..432c1ec4 100644 --- a/neo/process.go +++ b/neo/process.go @@ -1,15 +1,31 @@ package neo import ( + "fmt" + "strconv" + "github.com/gin-gonic/gin" "github.com/yaoapp/gou/process" "github.com/yaoapp/kun/exception" + "github.com/yaoapp/yao/neo/conversation" "github.com/yaoapp/yao/neo/message" ) +// GetNeo returns the Neo instance +func GetNeo() *DSL { + if Neo == nil { + exception.New("Neo is not initialized", 500).Throw() + } + return Neo +} + func init() { process.RegisterGroup("neo", map[string]process.Handler{ - "write": ProcessWrite, + "write": ProcessWrite, + "assistant.add": processAssistantAdd, + "assistant.save": processAssistantSave, + "assistant.delete": processAssistantDelete, + "assistant.search": processAssistantSearch, }) } @@ -39,3 +55,123 @@ func ProcessWrite(process *process.Process) interface{} { return nil } + +// processAssistantAdd process the assistant add request +func processAssistantAdd(process *process.Process) interface{} { + process.ValidateArgNums(1) + data := process.ArgsMap(0) + + neo := GetNeo() + if neo.Conversation == nil { + exception.New("Neo conversation is not initialized", 500).Throw() + } + + err := neo.Conversation.SaveAssistant(data) + if err != nil { + exception.New("Failed to add assistant: %s", 500, err.Error()).Throw() + } + + return data +} + +// processAssistantSave process the assistant save request +func processAssistantSave(process *process.Process) interface{} { + process.ValidateArgNums(1) + data := process.ArgsMap(0) + + neo := GetNeo() + if neo.Conversation == nil { + exception.New("Neo conversation is not initialized", 500).Throw() + } + + err := neo.Conversation.SaveAssistant(data) + if err != nil { + exception.New("Failed to save assistant: %s", 500, err.Error()).Throw() + } + + return data +} + +// processAssistantDelete process the assistant delete request +func processAssistantDelete(process *process.Process) interface{} { + process.ValidateArgNums(1) + assistantID := process.ArgsString(0) + + neo := GetNeo() + if neo.Conversation == nil { + exception.New("Neo conversation is not initialized", 500).Throw() + } + + err := neo.Conversation.DeleteAssistant(assistantID) + if err != nil { + exception.New("Failed to delete assistant: %s", 500, err.Error()).Throw() + } + + return gin.H{"message": "ok"} +} + +// processAssistantSearch process the assistant search request +func processAssistantSearch(process *process.Process) interface{} { + params := process.ArgsMap(0) + filter := conversation.AssistantFilter{} + + // Parse page and pagesize + if page, ok := params["page"]; ok { + pageStr := fmt.Sprintf("%v", page) + if pageInt, err := strconv.Atoi(pageStr); err == nil { + filter.Page = pageInt + } + } + if pagesize, ok := params["pagesize"]; ok { + pagesizeStr := fmt.Sprintf("%v", pagesize) + if pagesizeInt, err := strconv.Atoi(pagesizeStr); err == nil { + filter.PageSize = pagesizeInt + } + } + + // Parse tags + if tags, ok := params["tags"]; ok { + switch v := tags.(type) { + case []interface{}: + filter.Tags = make([]string, len(v)) + for i, tag := range v { + filter.Tags[i] = fmt.Sprintf("%v", tag) + } + case []string: + filter.Tags = v + } + } + + // Parse keywords + if keywords, ok := params["keywords"].(string); ok { + filter.Keywords = keywords + } + + // Parse connector + if connector, ok := params["connector"].(string); ok { + filter.Connector = connector + } + + // Parse mentionable + if mentionable, ok := params["mentionable"].(bool); ok { + filter.Mentionable = &mentionable + } + + // Parse automated + if automated, ok := params["automated"].(bool); ok { + filter.Automated = &automated + } + + // Get assistants + neo := GetNeo() + if neo.Conversation == nil { + exception.New("Neo conversation is not initialized", 500).Throw() + } + + res, err := neo.Conversation.GetAssistants(filter) + if err != nil { + exception.New("get assistants error: %s", 500, err).Throw() + } + + return res +} diff --git a/neo/process_test.go b/neo/process_test.go new file mode 100644 index 00000000..b7eca7f1 --- /dev/null +++ b/neo/process_test.go @@ -0,0 +1,390 @@ +package neo + +import ( + "fmt" + "testing" + + jsoniter "github.com/json-iterator/go" + "github.com/stretchr/testify/assert" + "github.com/yaoapp/gou/process" + "github.com/yaoapp/kun/any" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/test" +) + +func prepare(t *testing.T) { + test.Prepare(t, config.Conf) + err := Load(config.Conf) + if err != nil { + t.Fatal(err) + } + + // Clean up the test data before each test + p, err := process.Of("neo.assistant.search", map[string]interface{}{ + "page": 1, + "pagesize": 1000, // Use a large page size to get all records + }) + if err != nil { + t.Fatal(err) + } + output, err := p.Exec() + if err != nil { + t.Fatal(err) + } + res := any.Of(output).Map() + items := res.Get("data") + if items != nil { + for _, item := range items.([]map[string]interface{}) { + assistantID := item["assistant_id"].(string) + p, err = process.Of("neo.assistant.delete", assistantID) + if err != nil { + t.Fatal(err) + } + _, err = p.Exec() + if err != nil { + t.Fatal(err) + } + } + } + + // Verify cleanup + p, err = process.Of("neo.assistant.search") + if err != nil { + t.Fatal(err) + } + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + res = any.Of(output).Map() + total := res.Get("total") + if total != nil && any.Of(total).CInt() > 0 { + t.Fatalf("Failed to clean up test data, %d records remaining", any.Of(total).CInt()) + } + + check(t) +} + +func TestProcessAssistantCRUD(t *testing.T) { + prepare(t) + defer test.Clean() + + // Create an assistant + tagsJSON, err := jsoniter.MarshalToString([]string{"tag1", "tag2", "tag3"}) + if err != nil { + t.Fatal(err) + } + + optionsJSON, err := jsoniter.MarshalToString(map[string]interface{}{ + "model": "gpt-4", + }) + if err != nil { + t.Fatal(err) + } + + assistant := map[string]interface{}{ + "name": "Test Assistant", + "type": "assistant", + "avatar": "https://example.com/avatar.png", + "connector": "openai", + "description": "Test Description", + "tags": tagsJSON, + "options": optionsJSON, + "mentionable": true, + "automated": true, + } + + // Test processAssistantAdd + p, err := process.Of("neo.assistant.add", assistant) + if err != nil { + t.Fatal(err) + } + + output, err := p.Exec() + if err != nil { + t.Fatal(err) + } + + res := any.Of(output).Map() + assert.Equal(t, "Test Assistant", res.Get("name")) + assert.NotEmpty(t, res.Get("assistant_id")) + assistantID := res.Get("assistant_id").(string) + + // Test processAssistantSearch - no filter + p, err = process.Of("neo.assistant.search") + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + searchRes := any.Of(output).Map() + total := searchRes.Get("total") + if total == nil { + total = int64(0) + } + assert.Equal(t, int64(1), total) + + items := searchRes.Get("data") + if items == nil { + items = []map[string]interface{}{} + } + assert.Equal(t, 1, len(items.([]map[string]interface{}))) + + // Test processAssistantSearch - with filter + p, err = process.Of("neo.assistant.search", map[string]interface{}{ + "tags": []string{"tag1"}, + "page": 1, + "pagesize": 10, + }) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + searchRes = any.Of(output).Map() + total = searchRes.Get("total") + if total == nil { + total = int64(0) + } + assert.Equal(t, int64(1), total) + + items = searchRes.Get("data") + if items == nil { + items = []map[string]interface{}{} + } + assert.Equal(t, 1, len(items.([]map[string]interface{}))) + + // Test processAssistantSave (Update) + assistant["assistant_id"] = assistantID + assistant["name"] = "Updated Assistant" + p, err = process.Of("neo.assistant.save", assistant) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + res = any.Of(output).Map() + assert.Equal(t, "Updated Assistant", res.Get("name")) + + // Test processAssistantDelete + p, err = process.Of("neo.assistant.delete", assistantID) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + deleteRes := any.Of(output).Map() + assert.Equal(t, "ok", deleteRes.Get("message")) + + // Verify deletion with search + p, err = process.Of("neo.assistant.search") + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + searchRes = any.Of(output).Map() + total = searchRes.Get("total") + if total == nil { + total = int64(0) + } + assert.Equal(t, int64(0), total) + + items = searchRes.Get("data") + if items == nil { + items = []map[string]interface{}{} + } + assert.Equal(t, 0, len(items.([]map[string]interface{}))) +} + +func TestProcessAssistantSearchPagination(t *testing.T) { + prepare(t) + defer test.Clean() + + // Create multiple assistants for pagination testing + for i := 0; i < 25; i++ { + tagsJSON, err := jsoniter.MarshalToString([]string{fmt.Sprintf("tag%d", i%5)}) + if err != nil { + t.Fatal(err) + } + + assistant := map[string]interface{}{ + "name": fmt.Sprintf("Assistant %d", i), + "type": "assistant", + "connector": fmt.Sprintf("connector%d", i%3), + "description": fmt.Sprintf("Description %d", i), + "tags": tagsJSON, + "mentionable": i%2 == 0, + "automated": i%3 == 0, + } + + p, err := process.Of("neo.assistant.add", assistant) + if err != nil { + t.Fatal(err) + } + + _, err = p.Exec() + if err != nil { + t.Fatal(err) + } + } + + // Test first page + p, err := process.Of("neo.assistant.search", map[string]interface{}{ + "page": 1, + "pagesize": 10, + }) + if err != nil { + t.Fatal(err) + } + + output, err := p.Exec() + if err != nil { + t.Fatal(err) + } + + res := any.Of(output).Map() + total := res.Get("total") + if total == nil { + total = int64(0) + } + assert.Equal(t, int64(25), total) + + items := res.Get("data") + if items == nil { + items = []map[string]interface{}{} + } + assert.Equal(t, 10, len(items.([]map[string]interface{}))) + + pageCnt := res.Get("pagecnt") + if pageCnt == nil { + pageCnt = 1 + } + assert.Equal(t, 3, pageCnt) + + // Test second page + p, err = process.Of("neo.assistant.search", map[string]interface{}{ + "page": 2, + "pagesize": 10, + }) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + res = any.Of(output).Map() + items = res.Get("data") + if items == nil { + items = []map[string]interface{}{} + } + assert.Equal(t, 10, len(items.([]map[string]interface{}))) + + // Test last page + p, err = process.Of("neo.assistant.search", map[string]interface{}{ + "page": 3, + "pagesize": 10, + }) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + res = any.Of(output).Map() + items = res.Get("data") + if items == nil { + items = []map[string]interface{}{} + } + assert.Equal(t, 5, len(items.([]map[string]interface{}))) + + // Test filtering with tags + p, err = process.Of("neo.assistant.search", map[string]interface{}{ + "tags": []string{"tag0"}, + "page": 1, + "pagesize": 10, + }) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + res = any.Of(output).Map() + items = res.Get("data") + if items == nil { + items = []map[string]interface{}{} + } + assert.Equal(t, 5, len(items.([]map[string]interface{}))) +} + +func TestProcessAssistantValidation(t *testing.T) { + prepare(t) + defer test.Clean() + + // Test missing required fields + p, err := process.Of("neo.assistant.add", map[string]interface{}{}) + if err != nil { + t.Fatal(err) + } + + _, err = p.Exec() + assert.NotNil(t, err) + + // Test invalid assistant ID for delete + p, err = process.Of("neo.assistant.delete", "non-existent-id") + if err != nil { + t.Fatal(err) + } + + _, err = p.Exec() + assert.NotNil(t, err) + + // Test invalid page number + p, err = process.Of("neo.assistant.search", map[string]interface{}{ + "page": -1, + "pagesize": 10, + }) + if err != nil { + t.Fatal(err) + } + + output, err := p.Exec() + assert.Nil(t, err) + + res := any.Of(output).Map() + total := res.Get("total") + if total == nil { + total = int64(0) + } + assert.Equal(t, int64(0), total) +} diff --git a/test/utils.go b/test/utils.go index 066eb0eb..e9eaf427 100644 --- a/test/utils.go +++ b/test/utils.go @@ -4,6 +4,7 @@ import ( "fmt" "os" "path/filepath" + "regexp" "strings" "testing" "time" @@ -82,6 +83,50 @@ func Prepare(t *testing.T, cfg config.Config, rootEnv ...string) { // cfg.DataRoot = filepath.Join(root, "data") // } + var appData []byte + var appFile string + + // Read app setting + if has, _ := application.App.Exists("app.yao"); has { + appFile = "app.yao" + appData, err = application.App.Read("app.yao") + if err != nil { + t.Fatal(err) + } + + } else if has, _ := application.App.Exists("app.jsonc"); has { + appFile = "app.jsonc" + appData, err = application.App.Read("app.jsonc") + if err != nil { + t.Fatal(err) + } + + } else if has, _ := application.App.Exists("app.json"); has { + appFile = "app.json" + appData, err = application.App.Read("app.json") + if err != nil { + t.Fatal(err) + } + } else { + t.Fatal(fmt.Errorf("app.yao or app.jsonc or app.json does not exists")) + } + + // Replace $ENV with os.Getenv + var envRe = regexp.MustCompile(`\$ENV\.([0-9a-zA-Z_-]+)`) + appData = envRe.ReplaceAllFunc(appData, func(s []byte) []byte { + key := string(s[5:]) + val := os.Getenv(key) + if val == "" { + return s + } + return []byte(val) + }) + share.App = share.AppInfo{} + err = application.Parse(appFile, appData, &share.App) + if err != nil { + t.Fatal(err) + } + utils.Init() dbconnect(t, cfg) load(t, cfg)