From 375ad96e9da53e8f5f6300740f9e8e2f24a108a9 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 19 Dec 2024 15:13:38 +0800 Subject: [PATCH 01/17] Add generate endpoints and improve response handling in Neo API - Introduced new API endpoints for generating custom content, chat titles, and prompts, enhancing the functionality of the Neo API. - Implemented handleGenerateCustom, handleGenerateTitle, and handleGeneratePrompts methods to process generation requests with appropriate validation and error handling. - Refactored the GenerateWithAI method to support different content types and system prompts, improving the flexibility of AI-generated responses. - Updated the API router to include the new generation routes, ensuring seamless integration with existing middleware. --- neo/api.go | 167 ++++++++++++++++++++++++++++++++++++++++------- neo/neo.go | 36 +++++++--- neo/test_data.go | 30 +++++++++ 3 files changed, 203 insertions(+), 30 deletions(-) create mode 100644 neo/test_data.go diff --git a/neo/api.go b/neo/api.go index 2fd9289d..8f692db6 100644 --- a/neo/api.go +++ b/neo/api.go @@ -61,6 +61,11 @@ func (neo *DSL) API(router *gin.Engine, path string) error { // Mention api router.GET(path+"/mentions", append(middlewares, neo.handleMentions)...) + // Generate api + router.POST(path+"/generate", append(middlewares, neo.handleGenerateCustom)...) + router.POST(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...) + router.POST(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) + // Dangerous operations router.DELETE(path+"/dangerous/clear_chats", append(middlewares, neo.handleChatsDeleteAll)...) @@ -390,29 +395,8 @@ func (neo *DSL) handleMentions(c *gin.Context) { return } - // Add test data - testMentions := []Mention{ - { - ID: "assistant_1", - Name: "Alice AI", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Alice", - }, - { - ID: "assistant_2", - Name: "Bob Bot", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Bob", - }, - { - ID: "assistant_3", - Name: "Carol AI", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Carol", - }, - } - // Filter mentions by keywords + testMentions := TestMentions if keywords != "" { filtered := []Mention{} for _, m := range testMentions { @@ -534,3 +518,142 @@ func (neo *DSL) handleChatsDeleteAll(c *gin.Context) { c.JSON(200, gin.H{"message": "ok"}) c.Done() } + +// generateResponse is a helper struct to handle both SSE and HTTP responses +type generateResponse struct { + c *gin.Context + sid string + content string + result interface{} + err error +} + +// validate checks common validation rules +func (r *generateResponse) validate() bool { + if r.sid == "" { + r.c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + return false + } + + if r.content == "" { + r.c.JSON(400, gin.H{"message": "content is required", "code": 400}) + return false + } + + return true +} + +// send handles both SSE and HTTP responses +func (r *generateResponse) send(key string) { + if r.err != nil { + if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { + msg := message.New().Error(r.err.Error()).Done() + msg.Write(r.c.Writer) + } else { + r.c.JSON(500, gin.H{"message": r.err.Error(), "code": 500}) + } + return + } + + if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { + // Set headers for SSE + r.c.Header("Content-Type", "text/event-stream;charset=utf-8") + r.c.Header("Cache-Control", "no-cache") + r.c.Header("Connection", "keep-alive") + + msg := message.New().Map(gin.H{key: r.result}).Done() + msg.Write(r.c.Writer) + } else { + r.c.JSON(200, gin.H{key: r.result}) + } +} + +// handleGenerateTitle handles generating a chat title +func (neo *DSL) handleGenerateTitle(c *gin.Context) { + var body struct { + Content string `json:"content"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + + resp := &generateResponse{ + c: c, + sid: c.GetString("__sid"), + content: body.Content, + } + if !resp.validate() { + return + } + + ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + defer cancel() + + resp.result, resp.err = neo.GenerateChatTitle(ctx, resp.content, c) + resp.send("result") +} + +// handleGeneratePrompts handles generating prompts +func (neo *DSL) handleGeneratePrompts(c *gin.Context) { + var body struct { + Content string `json:"content"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + + resp := &generateResponse{ + c: c, + sid: c.GetString("__sid"), + content: body.Content, + } + if !resp.validate() { + return + } + + ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + defer cancel() + + resp.result, resp.err = neo.GeneratePrompts(ctx, resp.content, c) + resp.send("result") +} + +// handleGenerateCustom handles generating custom content +func (neo *DSL) handleGenerateCustom(c *gin.Context) { + var body struct { + Content string `json:"content"` + Type string `json:"type"` + SystemPrompt string `json:"system_prompt"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + + resp := &generateResponse{ + c: c, + sid: c.GetString("__sid"), + content: body.Content, + } + if !resp.validate() { + return + } + + // Additional validations for custom generation + if body.Type == "" { + c.JSON(400, gin.H{"message": "type is required", "code": 400}) + return + } + if body.SystemPrompt == "" { + c.JSON(400, gin.H{"message": "system_prompt is required", "code": 400}) + return + } + + ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + defer cancel() + + resp.result, resp.err = neo.GenerateWithAI(ctx, resp.content, body.Type, body.SystemPrompt, c) + resp.send("result") +} diff --git a/neo/neo.go b/neo/neo.go index 48ef866e..115ebc97 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -58,20 +58,40 @@ func (neo *DSL) GetMentions(keywords string) ([]Mention, error) { return neo.HookMention(context.Background(), keywords) } +// GeneratePrompts generate prompts for the AI assistant +func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context) (string, error) { + prompts := ` + Help me generate prompts for the AI assistant + 1. The prompts should guide the AI to better understand and respond to user questions + 2. The prompts should be clear and specific + 3. The prompts should be in the same language as the input + 4. Keep the prompts concise but comprehensive + ` + return neo.GenerateWithAI(ctx, input, "prompts", prompts, c) +} + // GenerateChatTitle generate the chat title func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (string, error) { - prompts := ` - Help me generate a title for the chat - 1. The title should be a short and concise description of the chat. - 2. The title should be a single sentence. - 3. The title should be in same language as the chat. - 4. The title should be no more than 50 characters. + Help me generate a title for the chat + 1. The title should be a short and concise description of the chat. + 2. The title should be a single sentence. + 3. The title should be in same language as the chat. + 4. The title should be no more than 50 characters. ` + return neo.GenerateWithAI(ctx, input, "title", prompts, c) +} +// GenerateWithAI generate content with AI, type can be "title", "prompts", etc. +func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context) (string, error) { messages := []map[string]interface{}{ - {"role": "system", "content": prompts}, - {"role": "user", "content": input}, + {"role": "system", "content": systemPrompt}, + { + "role": "user", + "content": input, + "type": messageType, + "name": ctx.Sid, + }, } res, err := neo.HookCreate(ctx, messages, c) diff --git a/neo/test_data.go b/neo/test_data.go new file mode 100644 index 00000000..f783b8de --- /dev/null +++ b/neo/test_data.go @@ -0,0 +1,30 @@ +package neo + +// TestMentions contains test data for mentions +var TestMentions = []Mention{ + { + ID: "assistant_1", + Name: "Alice AI", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Alice", + }, + { + ID: "assistant_2", + Name: "Bob Bot", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Bob", + }, + { + ID: "assistant_3", + Name: "Carol AI", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Carol", + }, + // ... 其他测试数据 ... + { + ID: "assistant_20", + Name: "Tara Tech", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Tara", + }, +} From fc1cbb2457db31e5921240aae3ca1ed91103854f Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 19 Dec 2024 15:16:08 +0800 Subject: [PATCH 02/17] Add OPTIONS routes for new generation endpoints in Neo API - Registered new OPTIONS routes for /generate, /generate/title, and /generate/prompts in the Neo API, enhancing CORS support for the recently introduced generation functionalities. - This update ensures that the API can handle preflight requests for these endpoints, improving overall API usability and compliance with CORS standards. --- neo/api.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/neo/api.go b/neo/api.go index 8f692db6..9d35b0c0 100644 --- a/neo/api.go +++ b/neo/api.go @@ -36,6 +36,9 @@ func (neo *DSL) API(router *gin.Engine, path string) error { router.OPTIONS(path+"/upload", neo.optionsHandler) router.OPTIONS(path+"/download", neo.optionsHandler) router.OPTIONS(path+"/mentions", neo.optionsHandler) + router.OPTIONS(path+"/generate", neo.optionsHandler) + router.OPTIONS(path+"/generate/title", neo.optionsHandler) + router.OPTIONS(path+"/generate/prompts", neo.optionsHandler) router.OPTIONS(path+"/dangerous/clear_chats", neo.optionsHandler) // Register endpoints with middlewares From a90d1d5a0aadf802bd5b4ac2b7a88c04b470e7f7 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 19 Dec 2024 16:27:51 +0800 Subject: [PATCH 03/17] Enhance Neo API with new generation endpoints and improved SSE handling - Added new GET and POST endpoints for generating custom content, chat titles, and prompts, expanding the API's capabilities. - Implemented support for Server-Sent Events (SSE) in response handling, allowing for real-time updates and improved user experience. - Refactored validation and error handling to differentiate between regular and SSE requests, ensuring appropriate responses based on the request type. - Updated the GenerateWithAI method to support silent mode for regular HTTP requests, enhancing flexibility in content generation. - Improved overall structure and maintainability of the API by consolidating request handling logic. --- neo/api.go | 142 ++++++++++++++++++++++++++++++++++++++++------------- neo/neo.go | 41 +++++++++++++--- 2 files changed, 141 insertions(+), 42 deletions(-) diff --git a/neo/api.go b/neo/api.go index 9d35b0c0..b680a711 100644 --- a/neo/api.go +++ b/neo/api.go @@ -65,8 +65,11 @@ func (neo *DSL) API(router *gin.Engine, path string) error { router.GET(path+"/mentions", append(middlewares, neo.handleMentions)...) // Generate api + router.GET(path+"/generate", append(middlewares, neo.handleGenerateCustom)...) router.POST(path+"/generate", append(middlewares, neo.handleGenerateCustom)...) + router.GET(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...) router.POST(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...) + router.GET(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) router.POST(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) // Dangerous operations @@ -534,12 +537,32 @@ type generateResponse struct { // validate checks common validation rules func (r *generateResponse) validate() bool { if r.sid == "" { - r.c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { + r.c.Header("Content-Type", "text/event-stream;charset=utf-8") + r.c.Header("Cache-Control", "no-cache") + r.c.Header("Connection", "keep-alive") + msg := message.New(). + Error("sid is required"). + Done() + msg.Write(r.c.Writer) + } else { + r.c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + } return false } if r.content == "" { - r.c.JSON(400, gin.H{"message": "content is required", "code": 400}) + if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { + r.c.Header("Content-Type", "text/event-stream;charset=utf-8") + r.c.Header("Cache-Control", "no-cache") + r.c.Header("Connection", "keep-alive") + msg := message.New(). + Error("content is required"). + Done() + msg.Write(r.c.Writer) + } else { + r.c.JSON(400, gin.H{"message": "content is required", "code": 400}) + } return false } @@ -550,7 +573,12 @@ func (r *generateResponse) validate() bool { func (r *generateResponse) send(key string) { if r.err != nil { if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { - msg := message.New().Error(r.err.Error()).Done() + r.c.Header("Content-Type", "text/event-stream;charset=utf-8") + r.c.Header("Cache-Control", "no-cache") + r.c.Header("Connection", "keep-alive") + msg := message.New(). + Error(r.err.Error()). + Done() msg.Write(r.c.Writer) } else { r.c.JSON(500, gin.H{"message": r.err.Error(), "code": 500}) @@ -559,12 +587,12 @@ func (r *generateResponse) send(key string) { } if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { - // Set headers for SSE r.c.Header("Content-Type", "text/event-stream;charset=utf-8") r.c.Header("Cache-Control", "no-cache") r.c.Header("Connection", "keep-alive") - - msg := message.New().Map(gin.H{key: r.result}).Done() + msg := message.New(). + Map(gin.H{key: r.result}). + Done() msg.Write(r.c.Writer) } else { r.c.JSON(200, gin.H{key: r.result}) @@ -573,18 +601,24 @@ func (r *generateResponse) send(key string) { // handleGenerateTitle handles generating a chat title func (neo *DSL) handleGenerateTitle(c *gin.Context) { - var body struct { - Content string `json:"content"` - } - if err := c.BindJSON(&body); err != nil { - c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) - return + var content string + if c.Request.Method == "GET" { + content = c.Query("content") + } else { + var body struct { + Content string `json:"content"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + content = body.Content } resp := &generateResponse{ c: c, sid: c.GetString("__sid"), - content: body.Content, + content: content, } if !resp.validate() { return @@ -593,25 +627,50 @@ func (neo *DSL) handleGenerateTitle(c *gin.Context) { ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") defer cancel() - resp.result, resp.err = neo.GenerateChatTitle(ctx, resp.content, c) + // Use silent mode for regular HTTP requests, streaming for SSE + silent := !strings.Contains(c.GetHeader("Accept"), "text/event-stream") + resp.result, resp.err = neo.GenerateChatTitle(ctx, resp.content, c, silent) resp.send("result") } // handleGeneratePrompts handles generating prompts func (neo *DSL) handleGeneratePrompts(c *gin.Context) { - var body struct { - Content string `json:"content"` - } - if err := c.BindJSON(&body); err != nil { - c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) - return + var content string + if c.Request.Method == "GET" { + content = c.Query("content") + } else { + var body struct { + Content string `json:"content"` + } + if err := c.BindJSON(&body); err != nil { + // For SSE requests, send error message in SSE format + if strings.Contains(c.GetHeader("Accept"), "text/event-stream") { + c.Header("Content-Type", "text/event-stream;charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + msg := message.New().Error("invalid request body").Done() + msg.Write(c.Writer) + return + } + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + content = body.Content } resp := &generateResponse{ c: c, sid: c.GetString("__sid"), - content: body.Content, + content: content, } + + // For SSE requests, set headers before validation + if strings.Contains(c.GetHeader("Accept"), "text/event-stream") { + c.Header("Content-Type", "text/event-stream;charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + } + if !resp.validate() { return } @@ -619,37 +678,50 @@ func (neo *DSL) handleGeneratePrompts(c *gin.Context) { ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") defer cancel() - resp.result, resp.err = neo.GeneratePrompts(ctx, resp.content, c) + // Use silent mode for regular HTTP requests, streaming for SSE + silent := !strings.Contains(c.GetHeader("Accept"), "text/event-stream") + resp.result, resp.err = neo.GeneratePrompts(ctx, resp.content, c, silent) resp.send("result") } // handleGenerateCustom handles generating custom content func (neo *DSL) handleGenerateCustom(c *gin.Context) { - var body struct { - Content string `json:"content"` - Type string `json:"type"` - SystemPrompt string `json:"system_prompt"` - } - if err := c.BindJSON(&body); err != nil { - c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) - return + var content, genType, systemPrompt string + + if c.Request.Method == "GET" { + content = c.Query("content") + genType = c.Query("type") + systemPrompt = c.Query("system_prompt") + } else { + var body struct { + Content string `json:"content"` + Type string `json:"type"` + SystemPrompt string `json:"system_prompt"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + content = body.Content + genType = body.Type + systemPrompt = body.SystemPrompt } resp := &generateResponse{ c: c, sid: c.GetString("__sid"), - content: body.Content, + content: content, } if !resp.validate() { return } // Additional validations for custom generation - if body.Type == "" { + if genType == "" { c.JSON(400, gin.H{"message": "type is required", "code": 400}) return } - if body.SystemPrompt == "" { + if systemPrompt == "" { c.JSON(400, gin.H{"message": "system_prompt is required", "code": 400}) return } @@ -657,6 +729,8 @@ func (neo *DSL) handleGenerateCustom(c *gin.Context) { ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") defer cancel() - resp.result, resp.err = neo.GenerateWithAI(ctx, resp.content, body.Type, body.SystemPrompt, c) + // Use silent mode for regular HTTP requests, streaming for SSE + silent := !strings.Contains(c.GetHeader("Accept"), "text/event-stream") + resp.result, resp.err = neo.GenerateWithAI(ctx, resp.content, genType, systemPrompt, c, silent) resp.send("result") } diff --git a/neo/neo.go b/neo/neo.go index 115ebc97..a83ec685 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -59,19 +59,25 @@ func (neo *DSL) GetMentions(keywords string) ([]Mention, error) { } // GeneratePrompts generate prompts for the AI assistant -func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context) (string, error) { +func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) { prompts := ` - Help me generate prompts for the AI assistant - 1. The prompts should guide the AI to better understand and respond to user questions + Optimize the prompts for the AI assistant + 1. Optimize prompts based on the user's input 2. The prompts should be clear and specific 3. The prompts should be in the same language as the input 4. Keep the prompts concise but comprehensive + 5. DO NOT ASK USER FOR MORE INFORMATION, JUST GENERATE PROMPTS + 6. DO NOT ANSWER THE QUESTION, JUST GENERATE PROMPTS ` - return neo.GenerateWithAI(ctx, input, "prompts", prompts, c) + isSilent := false + if len(silent) > 0 { + isSilent = silent[0] + } + return neo.GenerateWithAI(ctx, input, "prompts", prompts, c, isSilent) } // GenerateChatTitle generate the chat title -func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (string, error) { +func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) { prompts := ` Help me generate a title for the chat 1. The title should be a short and concise description of the chat. @@ -79,11 +85,15 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (st 3. The title should be in same language as the chat. 4. The title should be no more than 50 characters. ` - return neo.GenerateWithAI(ctx, input, "title", prompts, c) + isSilent := false + if len(silent) > 0 { + isSilent = silent[0] + } + return neo.GenerateWithAI(ctx, input, "title", prompts, c, isSilent) } // GenerateWithAI generate content with AI, type can be "title", "prompts", etc. -func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context) (string, error) { +func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) { messages := []map[string]interface{}{ {"role": "system", "content": systemPrompt}, { @@ -139,8 +149,21 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy // Append content and send message content = msg.Append(content) + // Only send real-time messages if not in silent mode + if !silent && msg.Message != nil && msg.Message.Text != "" { + message.New(). + Map(map[string]interface{}{ + "text": msg.Message.Text, + "done": msg.Message.Done, + }). + Write(c.Writer) + } + // Complete the stream if msg.Message.Done { + if !silent && msg.Message.Text == "" { + msg.Write(c.Writer) + } done <- true return 0 // break } @@ -151,7 +174,9 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy if err != nil { log.Error("Chat error: %s", err.Error()) - message.New().Error(err).Done().Write(c.Writer) + if !silent { + message.New().Error(err).Done().Write(c.Writer) + } } done <- true From c1e95374f2505606bd1030386f7871764866fb5f Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 19 Dec 2024 16:52:43 +0800 Subject: [PATCH 04/17] Update GenerateChatTitle method to include silent mode in chat handling - Modified the handleChatUpdate function to pass an additional parameter for silent mode when generating chat titles, enhancing the flexibility of chat title generation. - This change allows for improved handling of chat updates, ensuring that the API can better manage user expectations and interactions. --- neo/api.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neo/api.go b/neo/api.go index b680a711..e17aa8f9 100644 --- a/neo/api.go +++ b/neo/api.go @@ -452,7 +452,7 @@ func (neo *DSL) handleChatUpdate(c *gin.Context) { ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "") defer cancel() - title, err := neo.GenerateChatTitle(ctx, body.Content, c) + title, err := neo.GenerateChatTitle(ctx, body.Content, c, true) if err != nil { c.JSON(500, gin.H{"message": err.Error(), "code": 500}) c.Done() From 33b28b50f48d3632810fe6cdf443a614b71bce5a Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 19 Dec 2024 17:15:44 +0800 Subject: [PATCH 05/17] Enhance SSE error handling in handleGenerateTitle method - Added support for Server-Sent Events (SSE) in the handleGenerateTitle function to send error messages in the appropriate SSE format when invalid request bodies are encountered. - Updated response headers for SSE requests to ensure correct content type and connection settings, improving the API's handling of real-time updates and error reporting. --- neo/api.go | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/neo/api.go b/neo/api.go index e17aa8f9..163b3fad 100644 --- a/neo/api.go +++ b/neo/api.go @@ -609,6 +609,15 @@ func (neo *DSL) handleGenerateTitle(c *gin.Context) { Content string `json:"content"` } if err := c.BindJSON(&body); err != nil { + // For SSE requests, send error message in SSE format + if strings.Contains(c.GetHeader("Accept"), "text/event-stream") { + c.Header("Content-Type", "text/event-stream;charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + msg := message.New().Error("invalid request body").Done() + msg.Write(c.Writer) + return + } c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) return } @@ -620,6 +629,14 @@ func (neo *DSL) handleGenerateTitle(c *gin.Context) { sid: c.GetString("__sid"), content: content, } + + // For SSE requests, set headers before validation + if strings.Contains(c.GetHeader("Accept"), "text/event-stream") { + c.Header("Content-Type", "text/event-stream;charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + } + if !resp.validate() { return } From af52dd35a1a1b4d8a6d730f49760d82f6742bf75 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 20 Dec 2024 12:30:31 +0800 Subject: [PATCH 06/17] Refactor conversation history handling and enhance Assistant structure - Updated SaveHistory method across Mongo, Redis, and Weaviate to include a context parameter, improving flexibility in message storage. - Modified the Assistant struct in types.go to add new fields: Type, Avatar, and Flows, enhancing the representation of assistant attributes. - Adjusted the Xun implementation to support the new context handling in SaveHistory, ensuring messages can now include contextual information. - Updated tests to reflect changes in message structure and ensure proper functionality with the new context parameter. --- neo/assistant/types.go | 17 +-- neo/conversation/mongo.go | 2 +- neo/conversation/redis.go | 2 +- neo/conversation/types.go | 4 +- neo/conversation/weaviate.go | 2 +- neo/conversation/xun.go | 217 +++++++++++++++++++---------------- neo/conversation/xun_test.go | 48 ++------ neo/neo.go | 1 + 8 files changed, 140 insertions(+), 153 deletions(-) diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 03930582..35ad239d 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -30,13 +30,16 @@ type QueryParam struct { // Assistant the assistant type Assistant struct { - ID string `json:"assistant_id"` // Assistant ID - Name string `json:"name,omitempty"` // Assistant Name - Connector string `json:"connector"` // AI Connector - Description string `json:"description,omitempty"` // Assistant Description - Option map[string]interface{} `json:"option,omitempty"` // AI Option - Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts - API API `json:"-" yaml:"-"` // Assistant API + ID string `json:"assistant_id"` // Assistant ID + Type string `json:"type,omitempty"` // Assistant Type, default is assistant + Name string `json:"name,omitempty"` // Assistant Name + Avatar string `json:"avatar,omitempty"` // Assistant Avatar + Connector string `json:"connector"` // AI Connector + Description string `json:"description,omitempty"` // Assistant Description + Option map[string]interface{} `json:"option,omitempty"` // AI Option + Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts + Flows []map[string]interface{} `json:"flows,omitempty"` // Assistant Flows + API API `json:"-" yaml:"-"` // Assistant API } // File the file diff --git a/neo/conversation/mongo.go b/neo/conversation/mongo.go index 98c938ca..a30461fe 100644 --- a/neo/conversation/mongo.go +++ b/neo/conversation/mongo.go @@ -30,7 +30,7 @@ func (conv *Mongo) GetHistory(sid string, cid string) ([]map[string]interface{}, } // SaveHistory save the history -func (conv *Mongo) SaveHistory(sid string, messages []map[string]interface{}, cid string) error { +func (conv *Mongo) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { return nil } diff --git a/neo/conversation/redis.go b/neo/conversation/redis.go index afb1aa74..7da3f4d2 100644 --- a/neo/conversation/redis.go +++ b/neo/conversation/redis.go @@ -30,7 +30,7 @@ func (conv *Redis) GetHistory(sid string, cid string) ([]map[string]interface{}, } // SaveHistory save the history -func (conv *Redis) SaveHistory(sid string, messages []map[string]interface{}, cid string) error { +func (conv *Redis) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { return nil } diff --git a/neo/conversation/types.go b/neo/conversation/types.go index 205a99d6..63375e90 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -43,9 +43,7 @@ type Conversation interface { GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) GetChat(sid string, cid string) (*ChatInfo, error) GetHistory(sid string, cid string) ([]map[string]interface{}, error) - SaveHistory(sid string, messages []map[string]interface{}, cid string) error - GetRequest(sid string, rid string) ([]map[string]interface{}, error) - SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error + SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error DeleteChat(sid string, cid string) error DeleteAllChats(sid string) error UpdateChatTitle(sid string, cid string, title string) error diff --git a/neo/conversation/weaviate.go b/neo/conversation/weaviate.go index 7b251851..0131251e 100644 --- a/neo/conversation/weaviate.go +++ b/neo/conversation/weaviate.go @@ -30,7 +30,7 @@ func (conv *Weaviate) GetHistory(sid string, cid string) ([]map[string]interface } // SaveHistory save the history -func (conv *Weaviate) SaveHistory(sid string, messages []map[string]interface{}, cid string) error { +func (conv *Weaviate) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { return nil } diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 61a25b55..8ab43dc7 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -7,6 +7,7 @@ import ( "time" "github.com/google/uuid" + jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/session" "github.com/yaoapp/kun/log" @@ -23,13 +24,16 @@ type Xun struct { } type row struct { - Role string `json:"role"` - Name string `json:"name"` // User name - Content string `json:"content"` - Sid string `json:"sid"` - Rid string `json:"rid"` - Cid string `json:"cid"` // Chat ID from chat history - ExpiredAt interface{} `json:"expired_at"` + Role string `json:"role"` // Message role + Name string `json:"name"` // User name + Content string `json:"content"` // Message content + Sid string `json:"sid"` // Session ID + Cid string `json:"cid"` // Chat ID from chat history + UID string `json:"uid"` // User ID + Context map[string]interface{} `json:"context"` // Message context + CreatedAt time.Time `json:"created_at"` // Created time + UpdatedAt *time.Time `json:"updated_at"` // Updated time + ExpiredAt interface{} `json:"expired_at"` // Expired time } // Public interface methods and constructor remain exported: @@ -111,6 +115,11 @@ func (conv *Xun) initialize() error { return err } + // Initialize assistant table + if err := conv.initAssistantTable(); err != nil { + return err + } + return nil } @@ -126,11 +135,12 @@ func (conv *Xun) initHistoryTable() error { err = conv.schema.CreateTable(historyTable, func(table schema.Blueprint) { table.ID("id") table.String("sid", 255).Index() - table.String("rid", 255).Null().Index() table.String("cid", 200).Null().Index() + table.String("uid", 255).Null().Index() table.String("role", 200).Null().Index() table.String("name", 200).Null().Index() table.Text("content").Null() + table.JSON("context").Null() table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() table.TimestampTz("updated_at").Null().Index() table.TimestampTz("expired_at").Null().Index() @@ -148,7 +158,7 @@ func (conv *Xun) initHistoryTable() error { return err } - fields := []string{"id", "sid", "rid", "cid", "role", "name", "content", "created_at", "updated_at", "expired_at"} + fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"} for _, field := range fields { if !tab.HasColumn(field) { return fmt.Errorf("%s is required", field) @@ -198,6 +208,52 @@ func (conv *Xun) initChatTable() error { return nil } +func (conv *Xun) initAssistantTable() error { + assistantTable := conv.getAssistantTable() + has, err := conv.schema.HasTable(assistantTable) + if err != nil { + return err + } + + // Create the assistant table + if !has { + err = conv.schema.CreateTable(assistantTable, func(table schema.Blueprint) { + table.ID("id") + table.String("assistant_id", 200).Unique().Index() + table.String("type", 200).SetDefault("assistant").Index() // default is assistant + table.String("name", 200).Null() + table.String("avatar", 200).Null() + table.String("connector", 200).NotNull() + table.Text("description").Null() + table.JSON("option").Null() + table.JSON("prompts").Null() + table.JSON("flows").Null() + table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() + table.TimestampTz("updated_at").Null().Index() + }) + + if err != nil { + return err + } + log.Trace("Create the assistant table: %s", assistantTable) + } + + // Validate the table + tab, err := conv.schema.GetTable(assistantTable) + if err != nil { + return err + } + + fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "option", "prompts", "flows", "created_at", "updated_at"} + for _, field := range fields { + if !tab.HasColumn(field) { + return fmt.Errorf("%s is required", field) + } + } + + return nil +} + func (conv *Xun) getUserID(sid string) (string, error) { field := "user_id" if conv.setting.UserField != "" { @@ -217,13 +273,17 @@ func (conv *Xun) getUserID(sid string) (string, error) { } func (conv *Xun) getHistoryTable() string { - return conv.setting.Table + return conv.setting.Table + "_history" } func (conv *Xun) getChatTable() string { return conv.setting.Table + "_chat" } +func (conv *Xun) getAssistantTable() string { + return conv.setting.Table + "_assistant" +} + // UpdateChatTitle update the chat title func (conv *Xun) UpdateChatTitle(sid string, cid string, title string) error { userID, err := conv.getUserID(sid) @@ -380,7 +440,7 @@ func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, e } qb := conv.newQuery(). - Select("role", "name", "content"). + Select("role", "name", "content", "context", "uid", "created_at", "updated_at"). Where("sid", userID). Where("cid", cid). OrderBy("id", "desc") @@ -401,18 +461,23 @@ func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, e res := []map[string]interface{}{} for _, row := range rows { - res = append([]map[string]interface{}{{ - "role": row.Get("role"), - "name": row.Get("name"), - "content": row.Get("content"), - }}, res...) + message := map[string]interface{}{ + "role": row.Get("role"), + "name": row.Get("name"), + "content": row.Get("content"), + "context": row.Get("context"), + "uid": row.Get("uid"), + "created_at": row.Get("created_at"), + "updated_at": row.Get("updated_at"), + } + res = append([]map[string]interface{}{message}, res...) } return res, nil } // SaveHistory save the history -func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid string) error { +func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { if cid == "" { cid = uuid.New().String() // Generate a new UUID if cid is empty @@ -450,24 +515,49 @@ func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid // Save message history defer conv.clean() var expiredAt interface{} = nil - values := []row{} + values := []map[string]interface{}{} if conv.setting.TTL > 0 { expiredAt = time.Now().Add(time.Duration(conv.setting.TTL) * time.Second) } + now := time.Now() for _, message := range messages { - value := row{ - Role: message["role"].(string), - Name: "", - Content: message["content"].(string), - Sid: userID, - Cid: cid, - ExpiredAt: expiredAt, + // Type assertion safety checks + role, ok := message["role"].(string) + if !ok { + return fmt.Errorf("invalid role type in message: %v", message["role"]) } - if message["name"] != nil { - value.Name = message["name"].(string) + content, ok := message["content"].(string) + if !ok { + return fmt.Errorf("invalid content type in message: %v", message["content"]) } + + var contextRaw interface{} = nil + if context != nil { + contextRaw, err = jsoniter.MarshalToString(context) + if err != nil { + return err + } + } + + value := map[string]interface{}{ + "role": role, + "name": "", + "content": content, + "sid": userID, + "cid": cid, + "uid": userID, + "context": contextRaw, + "created_at": now, + "updated_at": nil, + "expired_at": expiredAt, + } + + if name, ok := message["name"].(string); ok { + value["name"] = name + } + values = append(values, value) } @@ -479,79 +569,6 @@ func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid return nil } -// GetRequest get the request history -func (conv *Xun) GetRequest(sid string, rid string) ([]map[string]interface{}, error) { - userID, err := conv.getUserID(sid) - if err != nil { - return nil, err - } - - qb := conv.newQuery(). - Select("role", "name", "content", "sid"). - Where("rid", rid). - Where("sid", userID). - OrderBy("id", "desc") - - if conv.setting.TTL > 0 { - qb.Where("expired_at", ">", time.Now()) - } - - limit := 20 - if conv.setting.MaxSize > 0 { - limit = conv.setting.MaxSize - } - - rows, err := qb.Limit(limit).Get() - if err != nil { - return nil, err - } - - res := []map[string]interface{}{} - for _, row := range rows { - res = append([]map[string]interface{}{{ - "role": row.Get("role"), - "name": row.Get("name"), - "content": row.Get("content"), - }}, res...) - } - - return res, nil -} - -// SaveRequest save the request history -func (conv *Xun) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error { - userID, err := conv.getUserID(sid) - if err != nil { - return err - } - - defer conv.clean() - var expiredAt interface{} = nil - values := []row{} - if conv.setting.TTL > 0 { - expiredAt = time.Now().Add(time.Duration(conv.setting.TTL) * time.Second) - } - - for _, message := range messages { - value := row{ - Role: message["role"].(string), - Name: "", - Content: message["content"].(string), - Sid: userID, - Cid: cid, - Rid: rid, - ExpiredAt: expiredAt, - } - - if message["name"] != nil { - value.Name = message["name"].(string) - } - values = append(values, value) - } - - return conv.newQuery().Insert(values) -} - // GetChat get the chat info and its history func (conv *Xun) GetChat(sid string, cid string) (*ChatInfo, error) { userID, err := conv.getUserID(sid) diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go index d42f32e2..34037b5c 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -45,7 +45,7 @@ func TestNewXunDefault(t *testing.T) { t.Fatal(err) } - fields := []string{"id", "sid", "cid", "rid", "role", "name", "content", "created_at", "updated_at", "expired_at"} + fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"} for _, field := range fields { assert.Equal(t, true, tab.HasColumn(field)) } @@ -103,7 +103,7 @@ func TestNewXunConnector(t *testing.T) { t.Fatal(err) } - fields := []string{"id", "sid", "cid", "rid", "role", "name", "content", "created_at", "updated_at", "expired_at"} + fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"} for _, field := range fields { assert.Equal(t, true, tab.HasColumn(field)) } @@ -143,7 +143,7 @@ func TestXunSaveAndGetHistory(t *testing.T) { err = conv.SaveHistory("123456", []map[string]interface{}{ {"role": "user", "name": "user1", "content": "hello"}, {"role": "assistant", "name": "user1", "content": "Hello there, how"}, - }, cid) + }, cid, nil) assert.Nil(t, err) // get the history @@ -154,38 +154,6 @@ func TestXunSaveAndGetHistory(t *testing.T) { assert.Equal(t, 2, len(data)) } -func TestXunSaveAndGetRequest(t *testing.T) { - - test.Prepare(t, config.Conf) - defer test.Clean() - defer capsule.Schema().DropTableIfExists("__unit_test_conversation") - - err := capsule.Schema().DropTableIfExists("__unit_test_conversation") - if err != nil { - t.Fatal(err) - } - - conv, err := NewXun(Setting{ - Connector: "default", - Table: "__unit_test_conversation", - TTL: 3600, - }) - - // save the history - err = conv.SaveRequest("123456", "912836", "test.command", []map[string]interface{}{ - {"role": "user", "name": "user1", "content": "hello"}, - {"role": "assistant", "name": "user1", "content": "Hello there, how"}, - }) - assert.Nil(t, err) - - // get the history - data, err := conv.GetRequest("123456", "912836") - if err != nil { - t.Fatal(err) - } - assert.Equal(t, 2, len(data)) -} - func TestXunSaveAndGetHistoryWithCID(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() @@ -209,7 +177,7 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) { {"role": "user", "name": "user1", "content": "hello"}, {"role": "assistant", "name": "assistant1", "content": "Hi! How can I help you?"}, } - err = conv.SaveHistory(sid, messages, cid) + err = conv.SaveHistory(sid, messages, cid, nil) assert.Nil(t, err) // get the history for specific cid @@ -224,7 +192,7 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) { moreMessages := []map[string]interface{}{ {"role": "user", "name": "user1", "content": "another message"}, } - err = conv.SaveHistory(sid, moreMessages, anotherCID) + err = conv.SaveHistory(sid, moreMessages, anotherCID, nil) assert.Nil(t, err) // get history for the first cid - should still be 2 messages @@ -294,7 +262,7 @@ func TestXunGetChats(t *testing.T) { } // Then save the history - err = conv.SaveHistory(sid, messages, chatID) + err = conv.SaveHistory(sid, messages, chatID, nil) if err != nil { t.Fatal(err) } @@ -344,7 +312,7 @@ func TestXunDeleteChat(t *testing.T) { } // Save the chat and history - err = conv.SaveHistory(sid, messages, cid) + err = conv.SaveHistory(sid, messages, cid, nil) assert.Nil(t, err) // Verify chat exists @@ -385,7 +353,7 @@ func TestXunDeleteAllChats(t *testing.T) { // Save multiple chats for i := 0; i < 3; i++ { cid := fmt.Sprintf("test_chat_%d", i) - err = conv.SaveHistory(sid, messages, cid) + err = conv.SaveHistory(sid, messages, cid, nil) assert.Nil(t, err) } diff --git a/neo/neo.go b/neo/neo.go index a83ec685..0583c1d9 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -498,6 +498,7 @@ func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages {"role": "assistant", "content": string(content), "name": sid}, }, chatID, + nil, ) if err != nil { From 76e7156daced50514c4a35d6eb6ee21ba9928a64 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 20 Dec 2024 15:48:28 +0800 Subject: [PATCH 07/17] Refactor assistant structure and replace base assistant with local implementation - Replaced the base assistant package with a local assistant implementation, enhancing modularity and maintainability. - Updated the newAssistantByConnector function to utilize the local assistant, ensuring consistent error handling and improved functionality. - Removed obsolete base assistant files, streamlining the codebase and reducing complexity. - Added a new 'mentionable' field to the assistant table in the Xun implementation, allowing for better management of assistant visibility in mentions. --- neo/assistant/{base => local}/chat.go | 4 ++-- neo/assistant/{base => local}/file.go | 10 +++++----- neo/assistant/{base/base.go => local/local.go} | 14 +++++++------- neo/conversation/xun.go | 16 ++-------------- neo/neo.go | 6 +++--- 5 files changed, 19 insertions(+), 31 deletions(-) rename neo/assistant/{base => local}/chat.go (65%) rename neo/assistant/{base => local}/file.go (89%) rename neo/assistant/{base/base.go => local/local.go} (63%) diff --git a/neo/assistant/base/chat.go b/neo/assistant/local/chat.go similarity index 65% rename from neo/assistant/base/chat.go rename to neo/assistant/local/chat.go index 7df5c58a..8c2518cb 100644 --- a/neo/assistant/base/chat.go +++ b/neo/assistant/local/chat.go @@ -1,4 +1,4 @@ -package base +package local import ( "context" @@ -6,7 +6,7 @@ import ( ) // Chat the chat -func (ast *Base) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error { +func (ast *Local) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error { if ast.openai == nil { return fmt.Errorf("api is not initialized") diff --git a/neo/assistant/base/file.go b/neo/assistant/local/file.go similarity index 89% rename from neo/assistant/base/file.go rename to neo/assistant/local/file.go index fd4f3029..79d40b6f 100644 --- a/neo/assistant/base/file.go +++ b/neo/assistant/local/file.go @@ -1,4 +1,4 @@ -package base +package local import ( "context" @@ -31,7 +31,7 @@ var AllowedFileTypes = map[string]string{ var MaxSize int64 = 20 * 1024 * 1024 // Upload the file -func (ast *Base) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*assistant.File, error) { +func (ast *Local) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*assistant.File, error) { // check file size if file.Size > MaxSize { @@ -69,13 +69,13 @@ func (ast *Base) Upload(ctx context.Context, file *multipart.FileHeader, reader }, nil } -func (ast *Base) id(temp string, ext string) (string, error) { +func (ast *Local) id(temp string, ext string) (string, error) { date := time.Now().Format("20060102") hash := fmt.Sprintf("%x", sha256.Sum256([]byte(temp)))[:8] return fmt.Sprintf("/__assistants/%s/%s/%s%s", ast.ID, date, hash, ext), nil } -func (ast *Base) allowed(contentType string) bool { +func (ast *Local) allowed(contentType string) bool { if _, ok := AllowedFileTypes[contentType]; ok { return true } @@ -87,7 +87,7 @@ func (ast *Base) allowed(contentType string) bool { } // Download downloads a file -func (ast *Base) Download(ctx context.Context, fileID string) (*assistant.FileResponse, error) { +func (ast *Local) Download(ctx context.Context, fileID string) (*assistant.FileResponse, error) { // Get the data filesystem data, err := fs.Get("data") diff --git a/neo/assistant/base/base.go b/neo/assistant/local/local.go similarity index 63% rename from neo/assistant/base/base.go rename to neo/assistant/local/local.go index bd608fe1..7a1386ff 100644 --- a/neo/assistant/base/base.go +++ b/neo/assistant/local/local.go @@ -1,4 +1,4 @@ -package base +package local import ( "context" @@ -8,16 +8,16 @@ import ( "github.com/yaoapp/yao/openai" ) -// Base the base assistant -type Base struct { +// Local the local assistant +type Local struct { ID string `json:"assistant_id"` Prompts []assistant.Prompt `json:"prompts,omitempty"` Connector connector.Connector `json:"-" yaml:"-"` openai *openai.OpenAI } -// New create a new base assistant -func New(connector connector.Connector, prompts []assistant.Prompt, id string) (*Base, error) { +// New create a new local assistant +func New(connector connector.Connector, prompts []assistant.Prompt, id string) (*Local, error) { setting := connector.Setting() api, err := openai.NewOpenAI(setting) @@ -25,10 +25,10 @@ func New(connector connector.Connector, prompts []assistant.Prompt, id string) ( return nil, err } - return &Base{Connector: connector, ID: id, Prompts: prompts, openai: api}, nil + return &Local{Connector: connector, ID: id, Prompts: prompts, openai: api}, nil } // List list all assistants -func (ast *Base) List(ctx context.Context, param assistant.QueryParam) ([]assistant.Assistant, error) { +func (ast *Local) List(ctx context.Context, param assistant.QueryParam) ([]assistant.Assistant, error) { return nil, nil } diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 8ab43dc7..5eecf4dc 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -23,19 +23,6 @@ type Xun struct { setting Setting } -type row struct { - Role string `json:"role"` // Message role - Name string `json:"name"` // User name - Content string `json:"content"` // Message content - Sid string `json:"sid"` // Session ID - Cid string `json:"cid"` // Chat ID from chat history - UID string `json:"uid"` // User ID - Context map[string]interface{} `json:"context"` // Message context - CreatedAt time.Time `json:"created_at"` // Created time - UpdatedAt *time.Time `json:"updated_at"` // Updated time - ExpiredAt interface{} `json:"expired_at"` // Expired time -} - // Public interface methods and constructor remain exported: // - NewXun // - UpdateChatTitle @@ -228,6 +215,7 @@ func (conv *Xun) initAssistantTable() error { table.JSON("option").Null() table.JSON("prompts").Null() table.JSON("flows").Null() + table.Boolean("mentionable").SetDefault(true).Index() // Whether this assistant can appear in @ mention list table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() table.TimestampTz("updated_at").Null().Index() }) @@ -244,7 +232,7 @@ func (conv *Xun) initAssistantTable() error { return err } - fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "option", "prompts", "flows", "created_at", "updated_at"} + fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "option", "prompts", "flows", "mentionable", "created_at", "updated_at"} for _, field := range fields { if !tab.HasColumn(field) { return fmt.Errorf("%s is required", field) diff --git a/neo/neo.go b/neo/neo.go index 0583c1d9..856082e6 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -11,7 +11,7 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/neo/assistant" - "github.com/yaoapp/yao/neo/assistant/base" + "github.com/yaoapp/yao/neo/assistant/local" "github.com/yaoapp/yao/neo/assistant/openai" "github.com/yaoapp/yao/neo/conversation" "github.com/yaoapp/yao/neo/message" @@ -417,9 +417,9 @@ func (neo *DSL) newAssistantByConnector(id string) (assistant.API, error) { } // Base on the assistant list hook - api, err := base.New(conn, neo.Prompts, id) + api, err := local.New(conn, neo.Prompts, id) if err != nil { - return nil, fmt.Errorf("Create base assistant error: %s", err.Error()) + return nil, fmt.Errorf("Create local assistant error: %s", err.Error()) } return api, nil } From 6f6c0b3d6ff5407689a92176dc0cdc40045ca656 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 20 Dec 2024 16:21:53 +0800 Subject: [PATCH 08/17] Enhance assistant table structure by adding 'files' field - Introduced a new 'files' field in the assistant table within the Xun implementation, allowing for better management of associated files. - Updated the fields array to include the new 'files' field, ensuring it is recognized as a required attribute in the assistant structure. --- neo/conversation/xun.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 5eecf4dc..b76fb66f 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -215,6 +215,7 @@ func (conv *Xun) initAssistantTable() error { table.JSON("option").Null() table.JSON("prompts").Null() table.JSON("flows").Null() + table.JSON("files").Null() table.Boolean("mentionable").SetDefault(true).Index() // Whether this assistant can appear in @ mention list table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() table.TimestampTz("updated_at").Null().Index() @@ -232,7 +233,7 @@ func (conv *Xun) initAssistantTable() error { return err } - fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "option", "prompts", "flows", "mentionable", "created_at", "updated_at"} + fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "option", "prompts", "flows", "files", "mentionable", "created_at", "updated_at"} for _, field := range fields { if !tab.HasColumn(field) { return fmt.Errorf("%s is required", field) From f022b79518a4071c07d96d5add8c92ad7160aaeb Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 20 Dec 2024 16:24:00 +0800 Subject: [PATCH 09/17] Add 'functions' field to assistant table in Xun implementation - Introduced a new 'functions' field in the assistant table, enhancing the structure for better management of assistant capabilities. - Updated the fields array to include 'functions', ensuring it is recognized as a required attribute in the assistant structure. --- neo/conversation/xun.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index b76fb66f..7ef29e4a 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -216,6 +216,7 @@ func (conv *Xun) initAssistantTable() error { table.JSON("prompts").Null() table.JSON("flows").Null() table.JSON("files").Null() + table.JSON("functions").Null() table.Boolean("mentionable").SetDefault(true).Index() // Whether this assistant can appear in @ mention list table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() table.TimestampTz("updated_at").Null().Index() @@ -233,7 +234,7 @@ func (conv *Xun) initAssistantTable() error { return err } - fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "option", "prompts", "flows", "files", "mentionable", "created_at", "updated_at"} + fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "option", "prompts", "flows", "files", "functions", "mentionable", "created_at", "updated_at"} for _, field := range fields { if !tab.HasColumn(field) { return fmt.Errorf("%s is required", field) From 8166503058b44bcd7544075eb89420cd55002202 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 28 Dec 2024 17:19:01 +0800 Subject: [PATCH 10/17] Refactor assistant table structure in Xun implementation - Enhanced the assistant table by adding new fields: 'options', 'permissions', 'readonly', and 'automated', improving the management of assistant attributes. - Updated the fields array to reflect the new structure, ensuring all fields are recognized as required attributes. - Improved code readability by adding comments to clarify the purpose of each field. --- neo/conversation/xun.go | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 7ef29e4a..3c394912 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -208,16 +208,19 @@ func (conv *Xun) initAssistantTable() error { table.ID("id") table.String("assistant_id", 200).Unique().Index() table.String("type", 200).SetDefault("assistant").Index() // default is assistant - table.String("name", 200).Null() - table.String("avatar", 200).Null() - table.String("connector", 200).NotNull() - table.Text("description").Null() - table.JSON("option").Null() - table.JSON("prompts").Null() - table.JSON("flows").Null() - table.JSON("files").Null() - table.JSON("functions").Null() - table.Boolean("mentionable").SetDefault(true).Index() // Whether this assistant can appear in @ mention list + table.String("name", 200).Null() // assistant name + table.String("avatar", 200).Null() // assistant avatar + table.String("connector", 200).NotNull() // assistant connector + table.Text("description").Null() // assistant description + table.JSON("options").Null() // assistant options + table.JSON("prompts").Null() // assistant prompts + table.JSON("flows").Null() // assistant flows + table.JSON("files").Null() // assistant files + table.JSON("functions").Null() // assistant functions + table.Boolean("readonly").SetDefault(false).Index() // assistant readonly + table.JSON("permissions").Null() // assistant permissions + table.Boolean("automated").SetDefault(true).Index() // assistant autoable + table.Boolean("mentionable").SetDefault(true).Index() // Whether this assistant can appear in @ mention list table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() table.TimestampTz("updated_at").Null().Index() }) @@ -234,7 +237,7 @@ func (conv *Xun) initAssistantTable() error { return err } - fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "option", "prompts", "flows", "files", "functions", "mentionable", "created_at", "updated_at"} + fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "mentionable", "created_at", "updated_at"} for _, field := range fields { if !tab.HasColumn(field) { return fmt.Errorf("%s is required", field) From b17778afea2c178e9377ee8633df0b85df4dada4 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 29 Dec 2024 10:03:38 +0800 Subject: [PATCH 11/17] Add assistant management functionality in conversation module - Introduced SaveAssistant, DeleteAssistant, and GetAssistants methods across Mongo, Redis, and Weaviate implementations to manage assistant records. - Enhanced the AssistantFilter and AssistantResponse types for improved pagination and filtering capabilities. - Updated the conversation interface to include new assistant-related methods, ensuring consistent functionality across different storage backends. - Added tests for assistant CRUD operations and pagination, ensuring robust functionality and reliability in managing assistants. --- neo/conversation/mongo.go | 28 ++++ neo/conversation/redis.go | 28 ++++ neo/conversation/types.go | 17 ++ neo/conversation/weaviate.go | 28 ++++ neo/conversation/xun.go | 83 ++++++++- neo/conversation/xun_test.go | 317 +++++++++++++++++++++++++++++++---- 6 files changed, 465 insertions(+), 36 deletions(-) diff --git a/neo/conversation/mongo.go b/neo/conversation/mongo.go index a30461fe..b406ffac 100644 --- a/neo/conversation/mongo.go +++ b/neo/conversation/mongo.go @@ -1,5 +1,7 @@ package conversation +import "github.com/yaoapp/xun" + // Mongo conversation type Mongo struct{} @@ -58,3 +60,29 @@ func (conv *Mongo) DeleteChat(sid string, cid string) error { func (conv *Mongo) DeleteAllChats(sid string) error { return nil } + +// SaveAssistant creates or updates an assistant +func (conv *Mongo) SaveAssistant(assistant map[string]interface{}) error { + return nil +} + +// DeleteAssistant deletes an assistant by assistant_id +func (conv *Mongo) DeleteAssistant(assistantID string) error { + return nil +} + +// 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, + }, + }, nil +} diff --git a/neo/conversation/redis.go b/neo/conversation/redis.go index 7da3f4d2..9aef5998 100644 --- a/neo/conversation/redis.go +++ b/neo/conversation/redis.go @@ -1,5 +1,7 @@ package conversation +import "github.com/yaoapp/xun" + // Redis conversation type Redis struct{} @@ -58,3 +60,29 @@ func (conv *Redis) DeleteChat(sid string, cid string) error { func (conv *Redis) DeleteAllChats(sid string) error { return nil } + +// SaveAssistant creates or updates an assistant +func (conv *Redis) SaveAssistant(assistant map[string]interface{}) error { + return nil +} + +// DeleteAssistant deletes an assistant by assistant_id +func (conv *Redis) DeleteAssistant(assistantID string) error { + return nil +} + +// 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, + }, + }, nil +} diff --git a/neo/conversation/types.go b/neo/conversation/types.go index 63375e90..fd9e2d40 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -1,5 +1,7 @@ package conversation +import "github.com/yaoapp/xun" + // Setting the conversation config type Setting struct { Connector string `json:"connector,omitempty"` @@ -38,6 +40,18 @@ type ChatGroupResponse struct { LastPage int `json:"last_page"` // 最后一页页码 } +// AssistantFilter represents the filter parameters for GetAssistants +type AssistantFilter struct { + Tags []string `json:"tags,omitempty"` + Page int `json:"page,omitempty"` // Page number, starting from 1 + PageSize int `json:"pagesize,omitempty"` // Items per page +} + +// AssistantResponse represents paginated assistant results +type AssistantResponse struct { + xun.P +} + // Conversation the store interface type Conversation interface { GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) @@ -47,4 +61,7 @@ type Conversation interface { DeleteChat(sid string, cid string) error DeleteAllChats(sid string) error UpdateChatTitle(sid string, cid string, title string) error + SaveAssistant(assistant map[string]interface{}) error + DeleteAssistant(assistantID string) error + GetAssistants(filter AssistantFilter) (*AssistantResponse, error) } diff --git a/neo/conversation/weaviate.go b/neo/conversation/weaviate.go index 0131251e..932aa22a 100644 --- a/neo/conversation/weaviate.go +++ b/neo/conversation/weaviate.go @@ -1,5 +1,7 @@ package conversation +import "github.com/yaoapp/xun" + // Weaviate Database conversation type Weaviate struct{} @@ -58,3 +60,29 @@ func (conv *Weaviate) DeleteChat(sid string, cid string) error { func (conv *Weaviate) DeleteAllChats(sid string) error { return nil } + +// SaveAssistant creates or updates an assistant +func (conv *Weaviate) SaveAssistant(assistant map[string]interface{}) error { + return nil +} + +// DeleteAssistant deletes an assistant by assistant_id +func (conv *Weaviate) DeleteAssistant(assistantID string) error { + return nil +} + +// 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, + }, + }, nil +} diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 3c394912..33eb6180 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -217,6 +217,7 @@ func (conv *Xun) initAssistantTable() error { table.JSON("flows").Null() // assistant flows table.JSON("files").Null() // assistant files table.JSON("functions").Null() // assistant functions + table.JSON("tags").Null() // assistant tags table.Boolean("readonly").SetDefault(false).Index() // assistant readonly table.JSON("permissions").Null() // assistant permissions table.Boolean("automated").SetDefault(true).Index() // assistant autoable @@ -237,7 +238,7 @@ func (conv *Xun) initAssistantTable() error { return err } - fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "mentionable", "created_at", "updated_at"} + fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "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) @@ -648,3 +649,83 @@ func (conv *Xun) DeleteAllChats(sid string) error { Delete() return err } + +// SaveAssistant creates or updates an assistant +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 + } + + // Check if assistant exists + exists, err := conv.query.New(). + Table(conv.getAssistantTable()). + Where("assistant_id", assistantID). + Exists() + if err != nil { + return err + } + + now := time.Now() + assistant["updated_at"] = now + + if exists { + // Update existing assistant + _, err = conv.query.New(). + Table(conv.getAssistantTable()). + Where("assistant_id", assistantID). + Update(assistant) + } else { + // Create new assistant + assistant["created_at"] = now + err = conv.query.New(). + Table(conv.getAssistantTable()). + Insert(assistant) + } + + return err +} + +// DeleteAssistant deletes an assistant by assistant_id +func (conv *Xun) DeleteAssistant(assistantID string) error { + _, err := conv.query.New(). + Table(conv.getAssistantTable()). + Where("assistant_id", assistantID). + Delete() + return err +} + +// GetAssistants retrieves assistants with pagination and tag filtering +func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { + qb := conv.query.New(). + Table(conv.getAssistantTable()) + + // Apply tag filter if provided + if filter.Tags != nil && len(filter.Tags) > 0 { + for i, tag := range filter.Tags { + if i == 0 { + qb.Where("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag)) + } else { + qb.OrWhere("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag)) + } + } + } + + // Set defaults for pagination + if filter.PageSize <= 0 { + filter.PageSize = 20 + } + if filter.Page <= 0 { + filter.Page = 1 + } + + // Get paginated results + paginator, err := qb.OrderBy("created_at", "desc"). + Paginate(filter.PageSize, filter.Page) + if err != nil { + return nil, err + } + + return &AssistantResponse{P: paginator}, nil +} diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go index 34037b5c..adc9306b 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -5,8 +5,10 @@ import ( "testing" "time" + 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" @@ -15,9 +17,21 @@ import ( func TestNewXunDefault(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") - err := capsule.Schema().DropTableIfExists("__unit_test_conversation") + err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") + if err != nil { + t.Fatal(err) + } + + err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") + if err != nil { + t.Fatal(err) + } + + err = capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") if err != nil { t.Fatal(err) } @@ -32,15 +46,29 @@ func TestNewXunDefault(t *testing.T) { return } - has, err := capsule.Schema().HasTable("__unit_test_conversation") + // Check history table + has, err := capsule.Schema().HasTable("__unit_test_conversation_history") if err != nil { t.Fatal(err) } - assert.Equal(t, true, has) - // validate the table - tab, err := conv.schema.GetTable(conv.setting.Table) + // Check chat table + has, err = capsule.Schema().HasTable("__unit_test_conversation_chat") + if err != nil { + t.Fatal(err) + } + assert.Equal(t, true, has) + + // Check assistant table + has, err = capsule.Schema().HasTable("__unit_test_conversation_assistant") + if err != nil { + t.Fatal(err) + } + assert.Equal(t, true, has) + + // validate the history table + tab, err := conv.schema.GetTable(conv.getHistoryTable()) if err != nil { t.Fatal(err) } @@ -50,17 +78,27 @@ func TestNewXunDefault(t *testing.T) { assert.Equal(t, true, tab.HasColumn(field)) } - conv, err = NewXun(Setting{ - Connector: "default", - Table: "__unit_test_conversation", - }) - - has, err = capsule.Schema().HasTable("__unit_test_conversation") + // validate the chat table + tab, err = conv.schema.GetTable(conv.getChatTable()) if err != nil { t.Fatal(err) } - assert.Equal(t, true, has) + chatFields := []string{"id", "chat_id", "title", "sid", "created_at", "updated_at"} + for _, field := range chatFields { + assert.Equal(t, true, tab.HasColumn(field)) + } + + // validate the assistant table + tab, err = conv.schema.GetTable(conv.getAssistantTable()) + if err != nil { + t.Fatal(err) + } + + assistantFields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "tags", "readonly", "permissions", "automated", "mentionable", "created_at", "updated_at"} + for _, field := range assistantFields { + assert.Equal(t, true, tab.HasColumn(field)) + } } func TestNewXunConnector(t *testing.T) { @@ -77,9 +115,14 @@ func TestNewXunConnector(t *testing.T) { t.Fatal(err) } - defer sch.DropTableIfExists("__unit_test_conversation") + defer sch.DropTableIfExists("__unit_test_conversation_history") + defer sch.DropTableIfExists("__unit_test_conversation_chat") + defer sch.DropTableIfExists("__unit_test_conversation_assistant") + + sch.DropTableIfExists("__unit_test_conversation_history") + sch.DropTableIfExists("__unit_test_conversation_chat") + sch.DropTableIfExists("__unit_test_conversation_assistant") - sch.DropTableIfExists("__unit_test_conversation") conv, err := NewXun(Setting{ Connector: "mysql", Table: "__unit_test_conversation", @@ -90,15 +133,29 @@ func TestNewXunConnector(t *testing.T) { return } - has, err := sch.HasTable("__unit_test_conversation") + // Check history table + has, err := sch.HasTable("__unit_test_conversation_history") if err != nil { t.Fatal(err) } - assert.Equal(t, true, has) - // validate the table - tab, err := conv.schema.GetTable(conv.setting.Table) + // Check chat table + has, err = sch.HasTable("__unit_test_conversation_chat") + if err != nil { + t.Fatal(err) + } + assert.Equal(t, true, has) + + // Check assistant table + has, err = sch.HasTable("__unit_test_conversation_assistant") + if err != nil { + t.Fatal(err) + } + assert.Equal(t, true, has) + + // validate the history table + tab, err := conv.schema.GetTable(conv.getHistoryTable()) if err != nil { t.Fatal(err) } @@ -108,26 +165,41 @@ func TestNewXunConnector(t *testing.T) { assert.Equal(t, true, tab.HasColumn(field)) } - conv, err = NewXun(Setting{ - Connector: "default", - Table: "__unit_test_conversation", - }) - - has, err = sch.HasTable("__unit_test_conversation") + // validate the chat table + tab, err = conv.schema.GetTable(conv.getChatTable()) if err != nil { t.Fatal(err) } - assert.Equal(t, true, has) + chatFields := []string{"id", "chat_id", "title", "sid", "created_at", "updated_at"} + for _, field := range chatFields { + assert.Equal(t, true, tab.HasColumn(field)) + } + + // validate the assistant table + tab, err = conv.schema.GetTable(conv.getAssistantTable()) + if err != nil { + t.Fatal(err) + } + + assistantFields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "tags", "readonly", "permissions", "automated", "mentionable", "created_at", "updated_at"} + for _, field := range assistantFields { + assert.Equal(t, true, tab.HasColumn(field)) + } } func TestXunSaveAndGetHistory(t *testing.T) { - test.Prepare(t, config.Conf) defer test.Clean() - defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") - err := capsule.Schema().DropTableIfExists("__unit_test_conversation") + err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") + if err != nil { + t.Fatal(err) + } + + err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") if err != nil { t.Fatal(err) } @@ -157,9 +229,15 @@ func TestXunSaveAndGetHistory(t *testing.T) { func TestXunSaveAndGetHistoryWithCID(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") - err := capsule.Schema().DropTableIfExists("__unit_test_conversation") + err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") + if err != nil { + t.Fatal(err) + } + + err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") if err != nil { t.Fatal(err) } @@ -220,11 +298,11 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) { func TestXunGetChats(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") // Drop both tables before test - err := capsule.Schema().DropTableIfExists("__unit_test_conversation") + err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") if err != nil { t.Fatal(err) } @@ -293,7 +371,7 @@ func TestXunGetChats(t *testing.T) { func TestXunDeleteChat(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") conv, err := NewXun(Setting{ @@ -333,7 +411,7 @@ func TestXunDeleteChat(t *testing.T) { func TestXunDeleteAllChats(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") conv, err := NewXun(Setting{ @@ -371,3 +449,172 @@ func TestXunDeleteAllChats(t *testing.T) { assert.Nil(t, err) assert.Equal(t, int64(0), response.Total) } + +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 + err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") + if err != nil { + t.Fatal(err) + } + + conv, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + if err != nil { + t.Fatal(err) + } + + // Test creating a new 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, + } + + // Test SaveAssistant (Create) + err = conv.SaveAssistant(assistant) + assert.Nil(t, err) + assistantID := assistant["assistant_id"].(string) + assert.NotEmpty(t, assistantID) + + // Test GetAssistants with no filter + resp, err := conv.GetAssistants(AssistantFilter{}) + assert.Nil(t, err) + assert.Equal(t, 1, len(resp.P.Items)) + + // 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)) + + // 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)) + + // 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)) + + // Test SaveAssistant (Update) + assistant["name"] = "Updated Assistant" + err = conv.SaveAssistant(assistant) + assert.Nil(t, err) + + 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, "Updated Assistant", item["name"]) + + // Test DeleteAssistant + err = conv.DeleteAssistant(assistantID) + assert.Nil(t, err) + + resp, err = conv.GetAssistants(AssistantFilter{}) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.P.Items)) +} + +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_assistant") + + // Drop assistant table before test + err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") + if err != nil { + t.Fatal(err) + } + + conv, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + if err != nil { + t.Fatal(err) + } + + // 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": "openai", + "description": fmt.Sprintf("Description %d", i), + "tags": tagsJSON, + } + err = conv.SaveAssistant(assistant) + assert.Nil(t, err) + } + + // Test first page + resp, err := conv.GetAssistants(AssistantFilter{ + Page: 1, + 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) + + // Test second page + resp, err = conv.GetAssistants(AssistantFilter{ + Page: 2, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Equal(t, 10, len(resp.P.Items)) + + // Test last page + resp, err = conv.GetAssistants(AssistantFilter{ + Page: 3, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Equal(t, 5, len(resp.P.Items)) + + // Test filtering with tags + resp, err = conv.GetAssistants(AssistantFilter{ + Tags: []string{"tag0"}, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Equal(t, 5, len(resp.P.Items)) +} From 8fafdc5855c52d10b95607859bb2b02d5c37df52 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 29 Dec 2024 11:06:41 +0800 Subject: [PATCH 12/17] Enhance assistant management in Neo API and conversation module - Added new endpoints for managing assistants: list, detail, save, and delete. - Implemented filtering capabilities for assistants based on keywords, tags, mentionable status, and automation. - Refactored the assistant structure to include new fields for better management and filtering. - Removed obsolete test data file, streamlining the codebase. - Updated tests to cover new assistant functionalities and filtering options, ensuring robust functionality. --- neo/api.go | 160 ++++++++++++++++++++++++++++++++--- neo/conversation/types.go | 118 +++++++++++++++++++------- neo/conversation/xun.go | 71 ++++++++++++---- neo/conversation/xun_test.go | 108 ++++++++++++++++++++++- neo/test_data.go | 30 ------- 5 files changed, 398 insertions(+), 89 deletions(-) delete mode 100644 neo/test_data.go diff --git a/neo/api.go b/neo/api.go index 163b3fad..cddab1e0 100644 --- a/neo/api.go +++ b/neo/api.go @@ -40,6 +40,8 @@ func (neo *DSL) API(router *gin.Engine, path string) error { router.OPTIONS(path+"/generate/title", neo.optionsHandler) router.OPTIONS(path+"/generate/prompts", neo.optionsHandler) router.OPTIONS(path+"/dangerous/clear_chats", neo.optionsHandler) + router.OPTIONS(path+"/assistants", neo.optionsHandler) + router.OPTIONS(path+"/assistants/:id", neo.optionsHandler) // Register endpoints with middlewares router.GET(path, append(middlewares, neo.handleChat)...) @@ -48,6 +50,12 @@ func (neo *DSL) API(router *gin.Engine, path string) error { // Status check router.GET(path+"/status", append(middlewares, neo.handleStatus)...) + // Assistant API + router.GET(path+"/assistants", append(middlewares, neo.handleAssistantList)...) + router.GET(path+"/assistants/:id", append(middlewares, neo.handleAssistantDetail)...) + router.POST(path+"/assistants", append(middlewares, neo.handleAssistantSave)...) + router.DELETE(path+"/assistants/:id", append(middlewares, neo.handleAssistantDelete)...) + // Chat api router.GET(path+"/chats", append(middlewares, neo.handleChatList)...) router.GET(path+"/chats/:id", append(middlewares, neo.handleChatDetail)...) @@ -394,28 +402,37 @@ func (neo *DSL) handleMentions(c *gin.Context) { // Get keywords from query parameter keywords := strings.ToLower(c.Query("keywords")) - mentions, err := neo.GetMentions(keywords) + mentionable := true + + // Query mentionable assistants + filter := conversation.AssistantFilter{ + Keywords: keywords, + Mentionable: &mentionable, + Page: 1, + PageSize: 20, + } + + response, err := neo.Conversation.GetAssistants(filter) if err != nil { c.JSON(500, gin.H{"message": err.Error(), "code": 500}) c.Done() return } - // Filter mentions by keywords - testMentions := TestMentions - if keywords != "" { - filtered := []Mention{} - for _, m := range testMentions { - if strings.Contains(strings.ToLower(m.Name), keywords) { - filtered = append(filtered, m) + // 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) } - testMentions = filtered } - // Append test data to actual mentions - mentions = append(mentions, testMentions...) - c.JSON(200, map[string]interface{}{"data": mentions}) c.Done() } @@ -751,3 +768,122 @@ func (neo *DSL) handleGenerateCustom(c *gin.Context) { resp.result, resp.err = neo.GenerateWithAI(ctx, resp.content, genType, systemPrompt, c, silent) resp.send("result") } + +// handleAssistantList handles listing assistants +func (neo *DSL) handleAssistantList(c *gin.Context) { + // Parse filter parameters + filter := conversation.AssistantFilter{ + Page: 1, + PageSize: 20, + } + + // Parse page and pagesize + if page := c.Query("page"); page != "" { + if n, err := strconv.Atoi(page); err == nil { + filter.Page = n + } + } + + if pageSize := c.Query("pagesize"); pageSize != "" { + if n, err := strconv.Atoi(pageSize); err == nil { + filter.PageSize = n + } + } + + // Parse tags + if tags := c.Query("tags"); tags != "" { + filter.Tags = strings.Split(tags, ",") + } + + response, err := neo.Conversation.GetAssistants(filter) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, map[string]interface{}{"data": response}) + c.Done() +} + +// handleAssistantDetail handles getting a single assistant's details +func (neo *DSL) handleAssistantDetail(c *gin.Context) { + assistantID := c.Param("id") + if assistantID == "" { + c.JSON(400, gin.H{"message": "assistant id is required", "code": 400}) + c.Done() + return + } + + filter := conversation.AssistantFilter{ + Page: 1, + PageSize: 1, + } + + response, err := neo.Conversation.GetAssistants(filter) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + // 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 + } + } + } + + if assistant == nil { + c.JSON(404, gin.H{"message": "assistant not found", "code": 404}) + c.Done() + return + } + + c.JSON(200, map[string]interface{}{"data": assistant}) + c.Done() +} + +// handleAssistantSave handles creating or updating an assistant +func (neo *DSL) handleAssistantSave(c *gin.Context) { + var assistant map[string]interface{} + if err := c.BindJSON(&assistant); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + c.Done() + return + } + + err := neo.Conversation.SaveAssistant(assistant) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, gin.H{"message": "ok", "data": assistant}) + c.Done() +} + +// handleAssistantDelete handles deleting an assistant +func (neo *DSL) handleAssistantDelete(c *gin.Context) { + assistantID := c.Param("id") + if assistantID == "" { + c.JSON(400, gin.H{"message": "assistant id is required", "code": 400}) + c.Done() + return + } + + err := neo.Conversation.DeleteAssistant(assistantID) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, gin.H{"message": "ok"}) + c.Done() +} diff --git a/neo/conversation/types.go b/neo/conversation/types.go index fd9e2d40..b85c8eeb 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -2,66 +2,126 @@ package conversation import "github.com/yaoapp/xun" -// Setting the conversation config +// Setting represents the conversation configuration structure +// Used to configure basic conversation parameters including connector, user field, table name, etc. type Setting struct { - Connector string `json:"connector,omitempty"` - UserField string `json:"user_field,omitempty"` // the user id field name, default is user_id - Table string `json:"table,omitempty"` - MaxSize int `json:"max_size,omitempty" yaml:"max_size,omitempty"` - TTL int `json:"ttl,omitempty" yaml:"ttl,omitempty"` + Connector string `json:"connector,omitempty"` // Name of the connector used to specify data storage method + UserField string `json:"user_field,omitempty"` // User ID field name, defaults to "user_id" + Table string `json:"table,omitempty"` // Database table name + MaxSize int `json:"max_size,omitempty" yaml:"max_size,omitempty"` // Maximum storage size limit + TTL int `json:"ttl,omitempty" yaml:"ttl,omitempty"` // Time To Live in seconds } -// ChatInfo represents the chat information and its history +// ChatInfo represents the chat information structure +// Contains basic information and history for a single chat type ChatInfo struct { - Chat map[string]interface{} `json:"chat"` - History []map[string]interface{} `json:"history"` + Chat map[string]interface{} `json:"chat"` // Basic chat information + History []map[string]interface{} `json:"history"` // Chat history records } -// ChatFilter represents the filter parameters for GetChats +// ChatFilter represents the chat filter structure +// Used for filtering and pagination when retrieving chat lists type ChatFilter struct { - Keywords string `json:"keywords,omitempty"` - Page int `json:"page,omitempty"` // 页码，从1开始 - PageSize int `json:"pagesize,omitempty"` // 每页数量 - Order string `json:"order,omitempty"` // desc/asc + Keywords string `json:"keywords,omitempty"` // Keyword search + Page int `json:"page,omitempty"` // Page number, starting from 1 + PageSize int `json:"pagesize,omitempty"` // Number of items per page + Order string `json:"order,omitempty"` // Sort order: desc/asc } -// ChatGroup represents a group of chats by date +// ChatGroup represents the chat group structure +// Groups chats by date type ChatGroup struct { - Label string `json:"label"` - Chats []map[string]interface{} `json:"chats"` + Label string `json:"label"` // Group label (typically a date) + Chats []map[string]interface{} `json:"chats"` // List of chats in this group } -// ChatGroupResponse represents paginated chat groups +// ChatGroupResponse represents the paginated chat group response +// Contains paginated chat group information type ChatGroupResponse struct { - Groups []ChatGroup `json:"groups"` - Page int `json:"page"` // 当前页码 - PageSize int `json:"pagesize"` // 每页数量 - Total int64 `json:"total"` // 总记录数 - LastPage int `json:"last_page"` // 最后一页页码 + Groups []ChatGroup `json:"groups"` // List of chat groups + Page int `json:"page"` // Current page number + PageSize int `json:"pagesize"` // Items per page + Total int64 `json:"total"` // Total number of records + LastPage int `json:"last_page"` // Last page number } -// AssistantFilter represents the filter parameters for GetAssistants +// AssistantFilter represents the assistant filter structure +// Used for filtering and pagination when retrieving assistant lists type AssistantFilter struct { - Tags []string `json:"tags,omitempty"` - Page int `json:"page,omitempty"` // Page number, starting from 1 - PageSize int `json:"pagesize,omitempty"` // Items per page + 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 + Mentionable *bool `json:"mentionable,omitempty"` // Filter by mentionable status + Automated *bool `json:"automated,omitempty"` // Filter by automation status + Page int `json:"page,omitempty"` // Page number, starting from 1 + PageSize int `json:"pagesize,omitempty"` // Items per page } -// AssistantResponse represents paginated assistant results +// AssistantResponse represents the assistant response structure +// Inherits from xun.P, used for returning paginated assistant lists type AssistantResponse struct { xun.P } -// Conversation the store interface +// Conversation defines the conversation storage interface +// Provides basic operations required for conversation management type Conversation interface { + // GetChats retrieves a list of chats + // sid: Session ID + // filter: Filter conditions + // Returns: Grouped chat list and potential error GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) + + // GetChat retrieves a single chat's information + // sid: Session ID + // cid: Chat ID + // Returns: Chat information and potential error GetChat(sid string, cid string) (*ChatInfo, error) + + // GetHistory retrieves chat history + // sid: Session ID + // cid: Chat ID + // Returns: History record list and potential error GetHistory(sid string, cid string) ([]map[string]interface{}, error) + + // SaveHistory saves chat history + // sid: Session ID + // messages: Message list + // cid: Chat ID + // context: Context information + // Returns: Potential error SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error + + // DeleteChat deletes a single chat + // sid: Session ID + // cid: Chat ID + // Returns: Potential error DeleteChat(sid string, cid string) error + + // DeleteAllChats deletes all chats + // sid: Session ID + // Returns: Potential error DeleteAllChats(sid string) error + + // UpdateChatTitle updates chat title + // sid: Session ID + // cid: Chat ID + // title: New title + // Returns: Potential error UpdateChatTitle(sid string, cid string, title string) error + + // SaveAssistant saves assistant information + // assistant: Assistant information + // Returns: Potential error SaveAssistant(assistant map[string]interface{}) error + + // DeleteAssistant deletes an assistant + // assistantID: Assistant ID + // Returns: Potential error DeleteAssistant(assistantID string) error + + // GetAssistants retrieves a list of assistants + // filter: Filter conditions + // Returns: Paginated assistant list and potential error GetAssistants(filter AssistantFilter) (*AssistantResponse, error) } diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 33eb6180..5e439e73 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -16,22 +16,34 @@ import ( "github.com/yaoapp/xun/dbal/schema" ) -// Xun Database conversation +// Package conversation provides functionality for managing chat conversations and assistants. + +// Xun implements the Conversation interface using a database backend. +// It provides functionality for: +// - Managing chat conversations and their message histories +// - Organizing chats with pagination and date-based grouping +// - Handling chat metadata like titles and creation dates +// - Managing AI assistants with their configurations and metadata +// - Supporting data expiration through TTL settings type Xun struct { query query.Query schema schema.Schema setting Setting } -// Public interface methods and constructor remain exported: -// - NewXun -// - UpdateChatTitle -// - GetChats -// - GetChat -// - GetHistory -// - SaveHistory -// - GetRequest -// - SaveRequest +// Public interface methods: +// +// NewXun creates a new conversation instance with the given settings +// UpdateChatTitle updates the title of a specific chat +// GetChats retrieves a paginated list of chats grouped by date +// GetChat retrieves a specific chat and its message history +// GetHistory retrieves the message history for a specific chat +// SaveHistory saves new messages to a chat's history +// DeleteChat deletes a specific chat and its history +// DeleteAllChats deletes all chats and their histories for a user +// SaveAssistant creates or updates an assistant +// DeleteAssistant deletes an assistant by assistant_id +// GetAssistants retrieves a paginated list of assistants with filtering // NewXun create a new conversation func NewXun(setting Setting) (*Xun, error) { @@ -696,20 +708,45 @@ func (conv *Xun) DeleteAssistant(assistantID string) error { return err } -// GetAssistants retrieves assistants with pagination and tag filtering +// GetAssistants retrieves assistants with pagination and filtering func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { qb := conv.query.New(). Table(conv.getAssistantTable()) // Apply tag filter if provided if filter.Tags != nil && len(filter.Tags) > 0 { - for i, tag := range filter.Tags { - if i == 0 { - qb.Where("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag)) - } else { - qb.OrWhere("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag)) + qb.Where(func(qb query.Query) { + for i, tag := range filter.Tags { + if i == 0 { + qb.Where("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag)) + } else { + qb.OrWhere("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag)) + } } - } + }) + } + + // 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 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) } // Set defaults for pagination diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go index adc9306b..90eacf74 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -483,6 +483,9 @@ func TestXunAssistantCRUD(t *testing.T) { t.Fatal(err) } + mentionable := true + automated := true + assistant := map[string]interface{}{ "name": "Test Assistant", "type": "assistant", @@ -491,6 +494,8 @@ func TestXunAssistantCRUD(t *testing.T) { "description": "Test Description", "tags": tagsJSON, "options": optionsJSON, + "mentionable": mentionable, + "automated": automated, } // Test SaveAssistant (Create) @@ -525,6 +530,45 @@ func TestXunAssistantCRUD(t *testing.T) { assert.Nil(t, err) assert.Equal(t, 0, len(resp.P.Items)) + // Test GetAssistants with keyword filter + resp, err = conv.GetAssistants(AssistantFilter{ + Keywords: "Test", + }) + assert.Nil(t, err) + assert.Equal(t, 1, len(resp.P.Items)) + + // Test GetAssistants with connector filter + resp, err = conv.GetAssistants(AssistantFilter{ + Connector: "openai", + }) + assert.Nil(t, err) + assert.Equal(t, 1, len(resp.P.Items)) + + // Test GetAssistants with mentionable filter + resp, err = conv.GetAssistants(AssistantFilter{ + Mentionable: &mentionable, + }) + assert.Nil(t, err) + assert.Equal(t, 1, len(resp.P.Items)) + + // Test GetAssistants with automated filter + resp, err = conv.GetAssistants(AssistantFilter{ + Automated: &automated, + }) + assert.Nil(t, err) + assert.Equal(t, 1, len(resp.P.Items)) + + // Test GetAssistants with combined filters + resp, err = conv.GetAssistants(AssistantFilter{ + Keywords: "Test", + Connector: "openai", + Mentionable: &mentionable, + Automated: &automated, + Tags: []string{"tag1"}, + }) + assert.Nil(t, err) + assert.Equal(t, 1, len(resp.P.Items)) + // Test SaveAssistant (Update) assistant["name"] = "Updated Assistant" err = conv.SaveAssistant(assistant) @@ -566,18 +610,30 @@ func TestXunAssistantPagination(t *testing.T) { } // Create multiple assistants for pagination testing + mentionable := true + automated := true 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), "type": "assistant", - "connector": "openai", + "connector": fmt.Sprintf("connector%d", i%3), "description": fmt.Sprintf("Description %d", i), "tags": tagsJSON, + "mentionable": mentionable, + "automated": automated, } err = conv.SaveAssistant(assistant) assert.Nil(t, err) @@ -617,4 +673,54 @@ func TestXunAssistantPagination(t *testing.T) { }) assert.Nil(t, err) assert.Equal(t, 5, len(resp.P.Items)) + + // Test filtering with keywords + resp, err = conv.GetAssistants(AssistantFilter{ + Keywords: "Assistant 1", + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Greater(t, len(resp.P.Items), 0) + + // Test filtering with connector + resp, err = conv.GetAssistants(AssistantFilter{ + Connector: "connector0", + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Greater(t, len(resp.P.Items), 0) + + // Test filtering with mentionable + mentionableTrue := true + resp, err = conv.GetAssistants(AssistantFilter{ + Mentionable: &mentionableTrue, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Greater(t, len(resp.P.Items), 0) + + // Test filtering with automated + automatedTrue := true + resp, err = conv.GetAssistants(AssistantFilter{ + Automated: &automatedTrue, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Greater(t, len(resp.P.Items), 0) + + // Test combined filters + resp, err = conv.GetAssistants(AssistantFilter{ + Tags: []string{"tag0"}, + Keywords: "Assistant", + Connector: "connector0", + Mentionable: &mentionableTrue, + Automated: &automatedTrue, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) } diff --git a/neo/test_data.go b/neo/test_data.go deleted file mode 100644 index f783b8de..00000000 --- a/neo/test_data.go +++ /dev/null @@ -1,30 +0,0 @@ -package neo - -// TestMentions contains test data for mentions -var TestMentions = []Mention{ - { - ID: "assistant_1", - Name: "Alice AI", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Alice", - }, - { - ID: "assistant_2", - Name: "Bob Bot", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Bob", - }, - { - ID: "assistant_3", - Name: "Carol AI", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Carol", - }, - // ... 其他测试数据 ... - { - ID: "assistant_20", - Name: "Tara Tech", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Tara", - }, -} From 697facb23d7111b7fcbb15613a085aa0fc71fa81 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 29 Dec 2024 11:13:57 +0800 Subject: [PATCH 13/17] Enhance Neo API with detailed endpoint examples and improved documentation - Added comprehensive curl examples for all API endpoints, including chat management, assistant management, file handling, and generation functionalities. - Improved comments throughout the code to clarify the purpose and usage of each endpoint, enhancing developer experience and usability. - Organized endpoint descriptions into categories for better readability and understanding of the API structure. --- neo/api.go | 79 ++++++++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 71 insertions(+), 8 deletions(-) diff --git a/neo/api.go b/neo/api.go index cddab1e0..c4a9ca5c 100644 --- a/neo/api.go +++ b/neo/api.go @@ -43,44 +43,107 @@ func (neo *DSL) API(router *gin.Engine, path string) error { router.OPTIONS(path+"/assistants", neo.optionsHandler) router.OPTIONS(path+"/assistants/:id", neo.optionsHandler) - // Register endpoints with middlewares + // Chat endpoint + // Example: + // curl -X GET 'http://localhost:5099/api/__yao/neo?content=Hello&chat_id=chat_123&context=previous_context&token=xxx' + // curl -X POST 'http://localhost:5099/api/__yao/neo' \ + // -H 'Content-Type: application/json' \ + // -d '{"content": "Hello", "chat_id": "chat_123", "context": "previous_context", "token": "xxx"}' router.GET(path, append(middlewares, neo.handleChat)...) router.POST(path, append(middlewares, neo.handleChat)...) - // Status check + // Status check endpoint + // Example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/status?token=xxx' router.GET(path+"/status", append(middlewares, neo.handleStatus)...) - // Assistant API + // Assistant API endpoints + // 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 assistant details example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/assistants/assistant_123?token=xxx' router.GET(path+"/assistants/:id", append(middlewares, neo.handleAssistantDetail)...) + + // Create/Update assistant example: + // curl -X POST 'http://localhost:5099/api/__yao/neo/assistants' \ + // -H 'Content-Type: application/json' \ + // -d '{"name": "My Assistant", "type": "chat", "tags": ["tag1", "tag2"], "mentionable": true, "avatar": "path/to/avatar.png", "token": "xxx"}' router.POST(path+"/assistants", append(middlewares, neo.handleAssistantSave)...) + + // Delete assistant example: + // curl -X DELETE 'http://localhost:5099/api/__yao/neo/assistants/assistant_123?token=xxx' router.DELETE(path+"/assistants/:id", append(middlewares, neo.handleAssistantDelete)...) - // Chat api + // Chat management endpoints + // List chats example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/chats?page=1&pagesize=20&keywords=search+term&order=desc&token=xxx' router.GET(path+"/chats", append(middlewares, neo.handleChatList)...) + + // Get chat details example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/chats/chat_123?token=xxx' router.GET(path+"/chats/:id", append(middlewares, neo.handleChatDetail)...) + + // Update chat example: + // curl -X POST 'http://localhost:5099/api/__yao/neo/chats/chat_123' \ + // -H 'Content-Type: application/json' \ + // -d '{"title": "New Title", "content": "Chat content for title generation", "token": "xxx"}' router.POST(path+"/chats/:id", append(middlewares, neo.handleChatUpdate)...) + + // Delete chat example: + // curl -X DELETE 'http://localhost:5099/api/__yao/neo/chats/chat_123?token=xxx' router.DELETE(path+"/chats/:id", append(middlewares, neo.handleChatDelete)...) - // History api + // Chat history endpoint + // Example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/history?chat_id=chat_123&token=xxx' router.GET(path+"/history", append(middlewares, neo.handleChatHistory)...) - // File api + // File management endpoints + // Upload file example: + // curl -X POST 'http://localhost:5099/api/__yao/neo/upload?chat_id=chat_123&token=xxx' \ + // -F 'file=@/path/to/file.txt' router.POST(path+"/upload", append(middlewares, neo.handleUpload)...) + + // Download file example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/download?file_id=file_123&disposition=attachment&token=xxx' \ + // -o downloaded_file.txt router.GET(path+"/download", append(middlewares, neo.handleDownload)...) - // Mention api + // Mentions endpoint + // Example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/mentions?keywords=assistant&token=xxx' router.GET(path+"/mentions", append(middlewares, neo.handleMentions)...) - // Generate api + // Generation endpoints + // Generate custom content example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/generate?content=Generate+something&type=custom&system_prompt=You+are+a+helpful+assistant&chat_id=chat_123&token=xxx' + // curl -X POST 'http://localhost:5099/api/__yao/neo/generate' \ + // -H 'Content-Type: application/json' \ + // -d '{"content": "Generate something", "type": "custom", "system_prompt": "You are a helpful assistant", "chat_id": "chat_123", "token": "xxx"}' router.GET(path+"/generate", append(middlewares, neo.handleGenerateCustom)...) router.POST(path+"/generate", append(middlewares, neo.handleGenerateCustom)...) + + // Generate title example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/generate/title?content=Chat+content&chat_id=chat_123&token=xxx' + // curl -X POST 'http://localhost:5099/api/__yao/neo/generate/title' \ + // -H 'Content-Type: application/json' \ + // -d '{"content": "Chat content", "chat_id": "chat_123", "token": "xxx"}' router.GET(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...) router.POST(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...) + + // Generate prompts example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/generate/prompts?content=Generate+prompts&chat_id=chat_123&token=xxx' + // curl -X POST 'http://localhost:5099/api/__yao/neo/generate/prompts' \ + // -H 'Content-Type: application/json' \ + // -d '{"content": "Generate prompts", "chat_id": "chat_123", "token": "xxx"}' router.GET(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) router.POST(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) // Dangerous operations + // Clear all chats example: + // curl -X DELETE 'http://localhost:5099/api/__yao/neo/dangerous/clear_chats?token=xxx' router.DELETE(path+"/dangerous/clear_chats", append(middlewares, neo.handleChatsDeleteAll)...) return nil From de49788d52d2a005192f5023e2b9ead7144f5245 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 29 Dec 2024 11:27:40 +0800 Subject: [PATCH 14/17] Add utility endpoint for listing connectors in Neo API - Introduced a new GET endpoint at `/utility/connectors` to list available connectors, enhancing the API's utility features. - Implemented `handleConnectors` function to filter and format connector data, providing a structured response with labels and values for each connector. - Updated documentation comments to include usage examples for the new endpoint, improving developer experience and usability. --- neo/api.go | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/neo/api.go b/neo/api.go index c4a9ca5c..0db39ee8 100644 --- a/neo/api.go +++ b/neo/api.go @@ -12,6 +12,7 @@ import ( "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/yaoapp/gou/api" + "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/process" "github.com/yaoapp/yao/helper" "github.com/yaoapp/yao/neo/conversation" @@ -141,6 +142,12 @@ func (neo *DSL) API(router *gin.Engine, path string) error { router.GET(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) router.POST(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) + // Utility endpoints + // List connectors example: + // curl -X GET 'http://localhost:5099/api/__yao/neo/utility/connectors?token=xxx' + router.GET(path+"/utility/connectors", append(middlewares, neo.handleConnectors)...) + + // Dangerous operations // Dangerous operations // Clear all chats example: // curl -X DELETE 'http://localhost:5099/api/__yao/neo/dangerous/clear_chats?token=xxx' @@ -950,3 +957,29 @@ func (neo *DSL) handleAssistantDelete(c *gin.Context) { c.JSON(200, gin.H{"message": "ok"}) c.Done() } + +// handleConnectors handles listing connectors +func (neo *DSL) handleConnectors(c *gin.Context) { + options := []map[string]interface{}{} + + // Filter and format connectors + for id, conn := range connector.Connectors { + if conn.Is(connector.OPENAI) || conn.Is(connector.MOAPI) { + setting := conn.Setting() + label := setting["label"] + if label == nil || label == "" { + label = setting["name"] + } + if label == nil || label == "" { + label = id + } + options = append(options, map[string]interface{}{ + "label": label, + "value": id, + }) + } + } + + c.JSON(200, gin.H{"data": options}) + c.Done() +} From af65a2d5fa6545014e3ef65152665c64c74c67a8 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 29 Dec 2024 12:57:17 +0800 Subject: [PATCH 15/17] 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) From 94f3ddb15825847d784fd20ca62d93e7213376a1 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 29 Dec 2024 17:14:01 +0800 Subject: [PATCH 16/17] Enhance assistant management and JSON field handling in Neo API - Updated SaveAssistant and processAssistantAdd methods to return assistant IDs, improving response consistency. - Implemented logic to handle JSON fields as both strings and native types, ensuring proper storage and retrieval of assistant attributes. - Refactored tests to cover various JSON field formats, including string, native types, and nil values, enhancing test coverage and reliability. - Improved the assistant management structure to support mixed JSON formats, ensuring robust functionality across different storage backends. --- neo/api.go | 7 +- neo/conversation/mongo.go | 80 +++++-------- neo/conversation/redis.go | 80 +++++-------- neo/conversation/types.go | 2 +- neo/conversation/weaviate.go | 81 +++++-------- neo/conversation/xun.go | 106 +++++++++++++---- neo/conversation/xun_test.go | 217 ++++++++++++++++++++-------------- neo/process.go | 8 +- neo/process_test.go | 218 ++++++++++++++++++++++++----------- 9 files changed, 463 insertions(+), 336 deletions(-) diff --git a/neo/api.go b/neo/api.go index 9da02e91..8ccf97ec 100644 --- a/neo/api.go +++ b/neo/api.go @@ -923,13 +923,18 @@ func (neo *DSL) handleAssistantSave(c *gin.Context) { return } - err := neo.Conversation.SaveAssistant(assistant) + id, err := neo.Conversation.SaveAssistant(assistant) if err != nil { c.JSON(500, gin.H{"message": err.Error(), "code": 500}) c.Done() return } + // Update the assistant map with the returned ID if it's not already set + if _, ok := assistant["assistant_id"]; !ok { + assistant["assistant_id"] = id + } + c.JSON(200, gin.H{"message": "ok", "data": assistant}) c.Done() } diff --git a/neo/conversation/mongo.go b/neo/conversation/mongo.go index c96fddf6..d4b20a16 100644 --- a/neo/conversation/mongo.go +++ b/neo/conversation/mongo.go @@ -1,83 +1,59 @@ package conversation -// Mongo conversation +// Mongo represents a MongoDB-based conversation storage type Mongo struct{} -// NewMongo create a new conversation +// NewMongo creates a new MongoDB conversation storage func NewMongo() *Mongo { return &Mongo{} } -// UpdateChatTitle update the chat title -func (conv *Mongo) UpdateChatTitle(sid string, cid string, title string) error { - return nil +// GetChats retrieves a list of chats +func (m *Mongo) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + return &ChatGroupResponse{}, nil } -// GetChats get the chat list -func (conv *Mongo) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { - return &ChatGroupResponse{ - Groups: []ChatGroup{}, - Page: filter.Page, - PageSize: filter.PageSize, - Total: 0, - LastPage: 1, - }, nil +// GetChat retrieves a single chat's information +func (m *Mongo) GetChat(sid string, cid string) (*ChatInfo, error) { + return &ChatInfo{}, nil } -// GetHistory get the history -func (conv *Mongo) GetHistory(sid string, cid string) ([]map[string]interface{}, error) { +// GetHistory retrieves chat history +func (m *Mongo) GetHistory(sid string, cid string) ([]map[string]interface{}, error) { return []map[string]interface{}{}, nil } -// SaveHistory save the history -func (conv *Mongo) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { +// SaveHistory saves chat history +func (m *Mongo) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { return nil } -// GetRequest get the request -func (conv *Mongo) GetRequest(sid string, rid string) ([]map[string]interface{}, error) { - return nil, nil -} - -// SaveRequest save the request -func (conv *Mongo) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error { +// DeleteChat deletes a single chat +func (m *Mongo) DeleteChat(sid string, cid string) error { return nil } -// GetChat get the chat info and its history -func (conv *Mongo) GetChat(sid string, cid string) (*ChatInfo, error) { - return nil, nil -} - -// DeleteChat deletes a specific chat and its history -func (conv *Mongo) DeleteChat(sid string, cid string) error { +// DeleteAllChats deletes all chats +func (m *Mongo) DeleteAllChats(sid string) error { return nil } -// DeleteAllChats deletes all chats and their histories for a user -func (conv *Mongo) DeleteAllChats(sid string) error { +// UpdateChatTitle updates chat title +func (m *Mongo) UpdateChatTitle(sid string, cid string, title string) error { return nil } -// SaveAssistant creates or updates an assistant -func (conv *Mongo) SaveAssistant(assistant map[string]interface{}) error { +// SaveAssistant saves assistant information +func (m *Mongo) SaveAssistant(assistant map[string]interface{}) (interface{}, error) { + return assistant["assistant_id"], nil +} + +// DeleteAssistant deletes an assistant +func (m *Mongo) DeleteAssistant(assistantID string) error { return nil } -// DeleteAssistant deletes an assistant by assistant_id -func (conv *Mongo) DeleteAssistant(assistantID string) error { - return nil -} - -// GetAssistants retrieves assistants with pagination and tag filtering -func (conv *Mongo) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { - return &AssistantResponse{ - Data: []map[string]interface{}{}, - Page: filter.Page, - PageSize: filter.PageSize, - PageCnt: 0, - Next: 0, - Prev: 0, - Total: 0, - }, nil +// GetAssistants retrieves a list of assistants +func (m *Mongo) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { + return &AssistantResponse{}, nil } diff --git a/neo/conversation/redis.go b/neo/conversation/redis.go index 7f88dbf5..fb150570 100644 --- a/neo/conversation/redis.go +++ b/neo/conversation/redis.go @@ -1,83 +1,59 @@ package conversation -// Redis conversation +// Redis represents a Redis-based conversation storage type Redis struct{} -// NewRedis create a new conversation +// NewRedis creates a new Redis conversation storage func NewRedis() *Redis { return &Redis{} } -// UpdateChatTitle update the chat title -func (conv *Redis) UpdateChatTitle(sid string, cid string, title string) error { - return nil +// GetChats retrieves a list of chats +func (r *Redis) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + return &ChatGroupResponse{}, nil } -// GetChats get the chat list -func (conv *Redis) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { - return &ChatGroupResponse{ - Groups: []ChatGroup{}, - Page: filter.Page, - PageSize: filter.PageSize, - Total: 0, - LastPage: 1, - }, nil +// GetChat retrieves a single chat's information +func (r *Redis) GetChat(sid string, cid string) (*ChatInfo, error) { + return &ChatInfo{}, nil } -// GetHistory get the history -func (conv *Redis) GetHistory(sid string, cid string) ([]map[string]interface{}, error) { +// GetHistory retrieves chat history +func (r *Redis) GetHistory(sid string, cid string) ([]map[string]interface{}, error) { return []map[string]interface{}{}, nil } -// SaveHistory save the history -func (conv *Redis) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { +// SaveHistory saves chat history +func (r *Redis) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { return nil } -// GetRequest get the request -func (conv *Redis) GetRequest(sid string, rid string) ([]map[string]interface{}, error) { - return nil, nil -} - -// SaveRequest save the request -func (conv *Redis) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error { +// DeleteChat deletes a single chat +func (r *Redis) DeleteChat(sid string, cid string) error { return nil } -// GetChat get the chat info and its history -func (conv *Redis) GetChat(sid string, cid string) (*ChatInfo, error) { - return nil, nil -} - -// DeleteChat deletes a specific chat and its history -func (conv *Redis) DeleteChat(sid string, cid string) error { +// DeleteAllChats deletes all chats +func (r *Redis) DeleteAllChats(sid string) error { return nil } -// DeleteAllChats deletes all chats and their histories for a user -func (conv *Redis) DeleteAllChats(sid string) error { +// UpdateChatTitle updates chat title +func (r *Redis) UpdateChatTitle(sid string, cid string, title string) error { return nil } -// SaveAssistant creates or updates an assistant -func (conv *Redis) SaveAssistant(assistant map[string]interface{}) error { +// SaveAssistant saves assistant information +func (r *Redis) SaveAssistant(assistant map[string]interface{}) (interface{}, error) { + return assistant["assistant_id"], nil +} + +// DeleteAssistant deletes an assistant +func (r *Redis) DeleteAssistant(assistantID string) error { return nil } -// DeleteAssistant deletes an assistant by assistant_id -func (conv *Redis) DeleteAssistant(assistantID string) error { - return nil -} - -// GetAssistants retrieves assistants with pagination and tag filtering -func (conv *Redis) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { - return &AssistantResponse{ - Data: []map[string]interface{}{}, - Page: filter.Page, - PageSize: filter.PageSize, - PageCnt: 0, - Next: 0, - Prev: 0, - Total: 0, - }, nil +// GetAssistants retrieves a list of assistants +func (r *Redis) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { + return &AssistantResponse{}, nil } diff --git a/neo/conversation/types.go b/neo/conversation/types.go index 87c0176a..c23fd83c 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -117,7 +117,7 @@ type Conversation interface { // SaveAssistant saves assistant information // assistant: Assistant information // Returns: Potential error - SaveAssistant(assistant map[string]interface{}) error + SaveAssistant(assistant map[string]interface{}) (interface{}, error) // DeleteAssistant deletes an assistant // assistantID: Assistant ID diff --git a/neo/conversation/weaviate.go b/neo/conversation/weaviate.go index 0b572111..8462e843 100644 --- a/neo/conversation/weaviate.go +++ b/neo/conversation/weaviate.go @@ -1,84 +1,59 @@ package conversation -// Weaviate Database conversation +// Weaviate represents a Weaviate-based conversation storage type Weaviate struct{} -// NewWeaviate create a new conversation +// NewWeaviate creates a new Weaviate conversation storage func NewWeaviate() *Weaviate { return &Weaviate{} } -// UpdateChatTitle update the chat title -func (conv *Weaviate) UpdateChatTitle(sid string, cid string, title string) error { - return nil +// GetChats retrieves a list of chats +func (w *Weaviate) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + return &ChatGroupResponse{}, nil } -// GetChats get the chat list -func (conv *Weaviate) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { - return &ChatGroupResponse{ - Groups: []ChatGroup{}, - Page: filter.Page, - PageSize: filter.PageSize, - Total: 0, - LastPage: 1, - }, nil +// GetChat retrieves a single chat's information +func (w *Weaviate) GetChat(sid string, cid string) (*ChatInfo, error) { + return &ChatInfo{}, nil } -// GetHistory get the history -func (conv *Weaviate) GetHistory(sid string, cid string) ([]map[string]interface{}, error) { +// GetHistory retrieves chat history +func (w *Weaviate) GetHistory(sid string, cid string) ([]map[string]interface{}, error) { return []map[string]interface{}{}, nil } -// SaveHistory save the history -func (conv *Weaviate) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { +// SaveHistory saves chat history +func (w *Weaviate) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { return nil } -// GetRequest get the request -func (conv *Weaviate) GetRequest(sid string, rid string) ([]map[string]interface{}, error) { - return nil, nil -} - -// SaveRequest save the request -func (conv *Weaviate) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error { +// DeleteChat deletes a single chat +func (w *Weaviate) DeleteChat(sid string, cid string) error { return nil } -// GetChat get the chat info and its history -func (conv *Weaviate) GetChat(sid string, cid string) (*ChatInfo, error) { - return nil, nil -} - -// DeleteChat deletes a specific chat and its history -func (conv *Weaviate) DeleteChat(sid string, cid string) error { +// DeleteAllChats deletes all chats +func (w *Weaviate) DeleteAllChats(sid string) error { return nil } -// DeleteAllChats deletes all chats and their histories for a user -func (conv *Weaviate) DeleteAllChats(sid string) error { +// UpdateChatTitle updates chat title +func (w *Weaviate) UpdateChatTitle(sid string, cid string, title string) error { return nil } -// SaveAssistant creates or updates an assistant -func (conv *Weaviate) SaveAssistant(assistant map[string]interface{}) error { +// SaveAssistant saves assistant information +func (w *Weaviate) SaveAssistant(assistant map[string]interface{}) (interface{}, error) { + return assistant["assistant_id"], nil +} + +// DeleteAssistant deletes an assistant +func (w *Weaviate) DeleteAssistant(assistantID string) error { return nil } -// DeleteAssistant deletes an assistant by assistant_id -func (conv *Weaviate) DeleteAssistant(assistantID string) error { - return nil -} - -// GetAssistants retrieves assistants with pagination and tag filtering -func (conv *Weaviate) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { - return &AssistantResponse{ - - Data: []map[string]interface{}{}, - Page: filter.Page, - PageSize: filter.PageSize, - PageCnt: 0, - Next: 0, - Prev: 0, - Total: 0, - }, nil +// GetAssistants retrieves a list of assistants +func (w *Weaviate) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { + return &AssistantResponse{}, nil } diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 752cda0e..6e017e90 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -662,53 +662,115 @@ func (conv *Xun) DeleteAllChats(sid string) error { return err } +// processJSONField processes a field that should be stored as JSON string +func (conv *Xun) processJSONField(field interface{}) (interface{}, error) { + if field == nil { + return nil, nil + } + + switch v := field.(type) { + case string: + return v, nil + default: + jsonStr, err := jsoniter.MarshalToString(v) + if err != nil { + return nil, fmt.Errorf("failed to marshal %v to JSON: %v", field, err) + } + return jsonStr, nil + } +} + +// parseJSONFields parses JSON string fields into their corresponding Go types +func (conv *Xun) parseJSONFields(data map[string]interface{}, fields []string) { + for _, field := range fields { + if val := data[field]; val != nil { + if strVal, ok := val.(string); ok && strVal != "" { + var parsed interface{} + if err := jsoniter.UnmarshalFromString(strVal, &parsed); err == nil { + data[field] = parsed + } + } + } + } +} + // SaveAssistant saves assistant information -func (conv *Xun) SaveAssistant(assistant map[string]interface{}) error { +func (conv *Xun) SaveAssistant(assistant map[string]interface{}) (interface{}, error) { // 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) + return nil, fmt.Errorf("field %s is required", field) } if assistant[field] == nil || assistant[field] == "" { - return fmt.Errorf("field %s cannot be empty", field) + return nil, 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) + // Create a copy of the assistant map to avoid modifying the original + assistantCopy := make(map[string]interface{}) + for k, v := range assistant { + assistantCopy[k] = v + } + + // Process JSON fields + jsonFields := []string{"tags", "options", "prompts", "flows", "files", "functions", "permissions"} + for _, field := range jsonFields { + if val, ok := assistantCopy[field]; ok && val != nil { + // If it's a string, try to parse it first + if strVal, ok := val.(string); ok && strVal != "" { + var parsed interface{} + if err := jsoniter.UnmarshalFromString(strVal, &parsed); err == nil { + assistantCopy[field] = parsed + } + } + } } // Generate assistant_id if not provided - if _, ok := assistant["assistant_id"]; !ok { - assistant["assistant_id"] = uuid.New().String() + if _, ok := assistantCopy["assistant_id"]; !ok { + assistantCopy["assistant_id"] = uuid.New().String() } // Check if assistant exists exists, err := conv.query.New(). Table(conv.getAssistantTable()). - Where("assistant_id", assistant["assistant_id"]). + Where("assistant_id", assistantCopy["assistant_id"]). Exists() if err != nil { - return err + return nil, err + } + + // Convert JSON fields to strings for storage + for _, field := range jsonFields { + if val, ok := assistantCopy[field]; ok && val != nil { + jsonStr, err := jsoniter.MarshalToString(val) + if err != nil { + return nil, fmt.Errorf("failed to marshal %s to JSON: %v", field, err) + } + assistantCopy[field] = jsonStr + } } // Update or insert if exists { - assistant["updated_at"] = time.Now() - _, err = conv.query.New(). + _, err := conv.query.New(). Table(conv.getAssistantTable()). - Where("assistant_id", assistant["assistant_id"]). - Update(assistant) - } else { - assistant["created_at"] = time.Now() - err = conv.query.New(). - Table(conv.getAssistantTable()). - Insert(assistant) + Where("assistant_id", assistantCopy["assistant_id"]). + Update(assistantCopy) + if err != nil { + return nil, err + } + return assistantCopy["assistant_id"], nil } - return err + err = conv.query.New(). + Table(conv.getAssistantTable()). + Insert(assistantCopy) + if err != nil { + return nil, err + } + return assistantCopy["assistant_id"], nil } // DeleteAssistant deletes an assistant by assistant_id @@ -812,10 +874,12 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro return nil, err } - // Convert rows to map slice + // Convert rows to map slice and parse JSON fields data := make([]map[string]interface{}, len(rows)) + jsonFields := []string{"tags", "options", "prompts", "flows", "files", "functions", "permissions"} for i, row := range rows { data[i] = row + conv.parseJSONFields(data[i], jsonFields) } return &AssistantResponse{ diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go index a08c800b..da3b4892 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -478,22 +478,10 @@ func TestXunAssistantCRUD(t *testing.T) { t.Fatal(err) } - // Test creating a new 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) - } - - mentionable := true - automated := true - + // Test creating a new assistant with different JSON field formats + // Test case 1: JSON fields as strings + tagsJSON := `["tag1", "tag2", "tag3"]` + optionsJSON := `{"model": "gpt-4"}` assistant := map[string]interface{}{ "name": "Test Assistant", "type": "assistant", @@ -502,95 +490,152 @@ func TestXunAssistantCRUD(t *testing.T) { "description": "Test Description", "tags": tagsJSON, "options": optionsJSON, - "mentionable": mentionable, - "automated": automated, + "mentionable": true, + "automated": true, } - // Test SaveAssistant (Create) - err = conv.SaveAssistant(assistant) + // Test SaveAssistant (Create) with string JSON + v, err := conv.SaveAssistant(assistant) assert.Nil(t, err) - assistantID := assistant["assistant_id"].(string) + assistantID := v.(string) assert.NotEmpty(t, assistantID) - // Test GetAssistants with no filter + // Test case 2: JSON fields as native types + assistant2 := map[string]interface{}{ + "name": "Test Assistant 2", + "type": "assistant", + "avatar": "https://example.com/avatar2.png", + "connector": "openai", + "description": "Test Description 2", + "tags": []string{"tag1", "tag2", "tag3"}, + "options": map[string]interface{}{"model": "gpt-4"}, + "prompts": []string{"prompt1", "prompt2"}, + "flows": []string{"flow1", "flow2"}, + "files": []string{"file1", "file2"}, + "functions": []map[string]interface{}{{"name": "func1"}, {"name": "func2"}}, + "permissions": map[string]interface{}{"read": true, "write": true}, + "mentionable": true, + "automated": true, + } + + // Test SaveAssistant (Create) with native types + v, err = conv.SaveAssistant(assistant2) + assert.Nil(t, err) + assistant2ID := v.(string) + assert.NotEmpty(t, assistant2ID) + + // Test case 3: Test with nil JSON fields + assistant3 := map[string]interface{}{ + "name": "Test Assistant 3", + "type": "assistant", + "connector": "openai", + "description": "Test Description 3", + "tags": nil, + "options": nil, + "prompts": nil, + "flows": nil, + "files": nil, + "functions": nil, + "permissions": nil, + "mentionable": true, + "automated": true, + } + + // Test SaveAssistant (Create) with nil fields + v, err = conv.SaveAssistant(assistant3) + assert.Nil(t, err) + assistant3ID := v.(string) + assert.NotEmpty(t, assistant3ID) + + // Test GetAssistants to verify JSON fields are properly stored resp, err := conv.GetAssistants(AssistantFilter{}) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.Data)) + assert.Equal(t, 3, 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.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) - // 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.Data)) + // 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"]) - // Test GetAssistants with non-existent tag - resp, err = conv.GetAssistants(AssistantFilter{ - Tags: []string{"nonexistent"}, - }) - assert.Nil(t, err) - assert.Equal(t, 0, len(resp.Data)) + // 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) - // Test GetAssistants with keyword filter - resp, err = conv.GetAssistants(AssistantFilter{ - Keywords: "Test", - }) - assert.Nil(t, err) - assert.Equal(t, 1, len(resp.Data)) + // 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"]) + break + } + } + assert.True(t, found) - // Test GetAssistants with connector filter - resp, err = conv.GetAssistants(AssistantFilter{ - Connector: "openai", - }) - assert.Nil(t, err) - 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.Data)) - - // Test GetAssistants with automated filter - resp, err = conv.GetAssistants(AssistantFilter{ - Automated: &automated, - }) - assert.Nil(t, err) - assert.Equal(t, 1, len(resp.Data)) - - // Test GetAssistants with combined filters - resp, err = conv.GetAssistants(AssistantFilter{ - Keywords: "Test", - Connector: "openai", - Mentionable: &mentionable, - Automated: &automated, - Tags: []string{"tag1"}, - }) - assert.Nil(t, err) - assert.Equal(t, 1, len(resp.Data)) - - // Test SaveAssistant (Update) - assistant["name"] = "Updated Assistant" - err = conv.SaveAssistant(assistant) + // Test updating with mixed JSON formats + assistant2["assistant_id"] = assistant2ID + _, err = conv.SaveAssistant(assistant2) assert.Nil(t, err) + // Verify update resp, err = conv.GetAssistants(AssistantFilter{}) assert.Nil(t, err) - assert.Equal(t, 1, len(resp.Data)) - item := resp.Data[0] - assert.Equal(t, "Updated Assistant", item["name"]) + 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 DeleteAssistant err = conv.DeleteAssistant(assistantID) assert.Nil(t, err) + err = conv.DeleteAssistant(assistant2ID) + assert.Nil(t, err) + err = conv.DeleteAssistant(assistant3ID) + assert.Nil(t, err) resp, err = conv.GetAssistants(AssistantFilter{}) assert.Nil(t, err) @@ -647,7 +692,7 @@ func TestXunAssistantPagination(t *testing.T) { "mentionable": mentionable, "automated": automated, } - err = conv.SaveAssistant(assistant) + _, err = conv.SaveAssistant(assistant) assert.Nil(t, err) } diff --git a/neo/process.go b/neo/process.go index 432c1ec4..ac631428 100644 --- a/neo/process.go +++ b/neo/process.go @@ -66,12 +66,12 @@ func processAssistantAdd(process *process.Process) interface{} { exception.New("Neo conversation is not initialized", 500).Throw() } - err := neo.Conversation.SaveAssistant(data) + id, err := neo.Conversation.SaveAssistant(data) if err != nil { exception.New("Failed to add assistant: %s", 500, err.Error()).Throw() } - return data + return id } // processAssistantSave process the assistant save request @@ -84,12 +84,12 @@ func processAssistantSave(process *process.Process) interface{} { exception.New("Neo conversation is not initialized", 500).Throw() } - err := neo.Conversation.SaveAssistant(data) + id, err := neo.Conversation.SaveAssistant(data) if err != nil { exception.New("Failed to save assistant: %s", 500, err.Error()).Throw() } - return data + return id } // processAssistantDelete process the assistant delete request diff --git a/neo/process_test.go b/neo/process_test.go index b7eca7f1..7dce8f1c 100644 --- a/neo/process_test.go +++ b/neo/process_test.go @@ -4,7 +4,6 @@ import ( "fmt" "testing" - jsoniter "github.com/json-iterator/go" "github.com/stretchr/testify/assert" "github.com/yaoapp/gou/process" "github.com/yaoapp/kun/any" @@ -69,19 +68,9 @@ 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) - } - + // Create an assistant with string JSON fields + tagsJSON := `["tag1", "tag2", "tag3"]` + optionsJSON := `{"model": "gpt-4"}` assistant := map[string]interface{}{ "name": "Test Assistant", "type": "assistant", @@ -94,7 +83,7 @@ func TestProcessAssistantCRUD(t *testing.T) { "automated": true, } - // Test processAssistantAdd + // Test processAssistantAdd with string JSON p, err := process.Of("neo.assistant.add", assistant) if err != nil { t.Fatal(err) @@ -105,12 +94,73 @@ func TestProcessAssistantCRUD(t *testing.T) { 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) + assistantID := output + assert.NotNil(t, assistantID) - // Test processAssistantSearch - no filter + // Test with native type JSON fields + assistant2 := map[string]interface{}{ + "name": "Test Assistant 2", + "type": "assistant", + "avatar": "https://example.com/avatar2.png", + "connector": "openai", + "description": "Test Description 2", + "tags": []string{"tag1", "tag2", "tag3"}, + "options": map[string]interface{}{"model": "gpt-4"}, + "prompts": []string{"prompt1", "prompt2"}, + "flows": []string{"flow1", "flow2"}, + "files": []string{"file1", "file2"}, + "functions": []map[string]interface{}{{"name": "func1"}, {"name": "func2"}}, + "permissions": map[string]interface{}{"read": true, "write": true}, + "mentionable": true, + "automated": true, + } + + // Test processAssistantAdd with native types + p, err = process.Of("neo.assistant.add", assistant2) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + assistant2ID := output + assert.NotNil(t, assistant2ID) + + // Test with nil JSON fields + assistant3 := map[string]interface{}{ + "name": "Test Assistant 3", + "type": "assistant", + "connector": "openai", + "description": "Test Description 3", + "tags": nil, + "options": nil, + "prompts": nil, + "flows": nil, + "files": nil, + "functions": nil, + "permissions": nil, + "mentionable": true, + "automated": true, + } + + // Test processAssistantAdd with nil fields + p, err = process.Of("neo.assistant.add", assistant3) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + assistant3ID := output + assert.NotNil(t, assistant3ID) + + // Test processAssistantSearch to verify all assistants p, err = process.Of("neo.assistant.search") if err != nil { t.Fatal(err) @@ -126,20 +176,68 @@ func TestProcessAssistantCRUD(t *testing.T) { if total == nil { total = int64(0) } - assert.Equal(t, int64(1), total) + assert.Equal(t, int64(3), total) items := searchRes.Get("data") if items == nil { items = []map[string]interface{}{} } - assert.Equal(t, 1, len(items.([]map[string]interface{}))) + assert.Equal(t, 3, 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, - }) + // Verify each assistant's JSON fields + for _, item := range items.([]map[string]interface{}) { + switch item["assistant_id"].(string) { + case assistantID: + assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"]) + assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"]) + case assistant2ID: + assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"]) + assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"]) + 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"]) + case assistant3ID: + 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"]) + } + } + + // Test updating with mixed JSON formats + assistant2["assistant_id"] = assistant2ID + assistant2["tags"] = `["tag4", "tag5"]` + assistant2["options"] = map[string]interface{}{"model": "gpt-3.5"} + p, err = process.Of("neo.assistant.save", assistant2) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + savedID := output + assert.NotNil(t, savedID) + + // Double check with a new search + p, err = process.Of("neo.assistant.search") if err != nil { t.Fatal(err) } @@ -150,33 +248,17 @@ func TestProcessAssistantCRUD(t *testing.T) { } 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{}{} + found := false + for _, item := range items.([]map[string]interface{}) { + if item["assistant_id"].(string) == assistant2ID { + found = true + assert.Equal(t, []interface{}{"tag4", "tag5"}, item["tags"]) + assert.Equal(t, map[string]interface{}{"model": "gpt-3.5"}, item["options"]) + break + } } - 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")) + assert.True(t, found) // Test processAssistantDelete p, err = process.Of("neo.assistant.delete", assistantID) @@ -192,7 +274,22 @@ func TestProcessAssistantCRUD(t *testing.T) { deleteRes := any.Of(output).Map() assert.Equal(t, "ok", deleteRes.Get("message")) - // Verify deletion with search + // Delete remaining assistants + p, err = process.Of("neo.assistant.delete", assistant2ID) + if err != nil { + t.Fatal(err) + } + _, err = p.Exec() + assert.Nil(t, err) + + p, err = process.Of("neo.assistant.delete", assistant3ID) + if err != nil { + t.Fatal(err) + } + _, err = p.Exec() + assert.Nil(t, err) + + // Verify all assistants are deleted p, err = process.Of("neo.assistant.search") if err != nil { t.Fatal(err) @@ -209,12 +306,6 @@ func TestProcessAssistantCRUD(t *testing.T) { 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) { @@ -223,17 +314,12 @@ func TestProcessAssistantSearchPagination(t *testing.T) { // 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, + "tags": []string{fmt.Sprintf("tag%d", i%5)}, "mentionable": i%2 == 0, "automated": i%3 == 0, } From 882277d974a6779628462bc61a242748d709d092 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 29 Dec 2024 17:30:33 +0800 Subject: [PATCH 17/17] Enhance assistant filtering and management in Neo API - Added support for filtering assistants by keywords, connector, and mentionable status in the handleAssistantList function. - Introduced a new parseBoolValue utility function to handle various boolean formats for filtering. - Updated the AssistantFilter structure to include an AssistantID field and a Select field for specifying returned fields. - Refactored process functions to replace 'add' with 'create' for consistency and clarity in assistant management. - Implemented new processAssistantFind function to retrieve assistants by ID, improving data retrieval capabilities. - Enhanced tests to cover new filtering options and ensure robust functionality across assistant management operations. --- neo/api.go | 65 +++++++++++++++++++++------ neo/conversation/types.go | 16 ++++--- neo/conversation/xun.go | 33 +++++++++++++- neo/conversation/xun_test.go | 87 ++++++++++++++++++++++++++++++++++++ neo/process.go | 37 +++++++++++++-- neo/process_test.go | 53 ++++++++++++++++++---- 6 files changed, 258 insertions(+), 33 deletions(-) diff --git a/neo/api.go b/neo/api.go index 8ccf97ec..be5640a4 100644 --- a/neo/api.go +++ b/neo/api.go @@ -863,6 +863,37 @@ func (neo *DSL) handleAssistantList(c *gin.Context) { filter.Tags = strings.Split(tags, ",") } + // Parse keywords + if keywords := c.Query("keywords"); keywords != "" { + filter.Keywords = keywords + } + + // Parse connector + if connector := c.Query("connector"); connector != "" { + filter.Connector = connector + } + + // Parse select fields + if selectFields := c.Query("select"); selectFields != "" { + filter.Select = strings.Split(selectFields, ",") + } + + // Parse mentionable (support various boolean formats) + if mentionable := c.Query("mentionable"); mentionable != "" { + val := parseBoolValue(mentionable) + if val != nil { + filter.Mentionable = val + } + } + + // Parse automated (support various boolean formats) + if automated := c.Query("automated"); automated != "" { + val := parseBoolValue(automated) + if val != nil { + filter.Automated = val + } + } + response, err := neo.Conversation.GetAssistants(filter) if err != nil { c.JSON(500, gin.H{"message": err.Error(), "code": 500}) @@ -874,6 +905,22 @@ func (neo *DSL) handleAssistantList(c *gin.Context) { c.Done() } +// parseBoolValue parses various string formats into a boolean pointer +// Supports: 1, 0, "1", "0", "true", "false", etc. +func parseBoolValue(value string) *bool { + value = strings.ToLower(strings.TrimSpace(value)) + switch value { + case "1", "true", "yes", "on": + v := true + return &v + case "0", "false", "no", "off": + v := false + return &v + default: + return nil + } +} + // handleAssistantDetail handles getting a single assistant's details func (neo *DSL) handleAssistantDetail(c *gin.Context) { assistantID := c.Param("id") @@ -884,8 +931,9 @@ func (neo *DSL) handleAssistantDetail(c *gin.Context) { } filter := conversation.AssistantFilter{ - Page: 1, - PageSize: 1, + AssistantID: assistantID, + Page: 1, + PageSize: 1, } response, err := neo.Conversation.GetAssistants(filter) @@ -895,22 +943,13 @@ func (neo *DSL) handleAssistantDetail(c *gin.Context) { return } - // Find the assistant by ID - var assistant map[string]interface{} - for _, item := range response.Data { - if id, ok := item["id"].(string); ok && id == assistantID { - assistant = item - break - } - } - - if assistant == nil { + if len(response.Data) == 0 { c.JSON(404, gin.H{"message": "assistant not found", "code": 404}) c.Done() return } - c.JSON(200, map[string]interface{}{"data": assistant}) + c.JSON(200, map[string]interface{}{"data": response.Data[0]}) c.Done() } diff --git a/neo/conversation/types.go b/neo/conversation/types.go index c23fd83c..521fc176 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -46,13 +46,15 @@ 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 - Mentionable *bool `json:"mentionable,omitempty"` // Filter by mentionable status - Automated *bool `json:"automated,omitempty"` // Filter by automation status - Page int `json:"page,omitempty"` // Page number, starting from 1 - PageSize int `json:"pagesize,omitempty"` // Items per page + 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 + 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/conversation/xun.go b/neo/conversation/xun.go index 6e017e90..d6ee4efa 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -829,6 +829,11 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro 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) @@ -865,6 +870,15 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro prevPage = 0 } + // Apply select fields if provided + if filter.Select != nil && len(filter.Select) > 0 { + selectFields := make([]interface{}, len(filter.Select)) + for i, field := range filter.Select { + selectFields[i] = field + } + qb.Select(selectFields...) + } + // Get paginated results rows, err := qb.OrderBy("created_at", "desc"). Offset(offset). @@ -879,7 +893,24 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro jsonFields := []string{"tags", "options", "prompts", "flows", "files", "functions", "permissions"} for i, row := range rows { data[i] = row - conv.parseJSONFields(data[i], jsonFields) + // Only parse JSON fields if they are selected or no select filter is provided + if filter.Select == nil || len(filter.Select) == 0 { + conv.parseJSONFields(data[i], jsonFields) + } else { + // Parse only selected JSON fields + selectedJSONFields := []string{} + for _, field := range jsonFields { + for _, selected := range filter.Select { + if selected == field { + selectedJSONFields = append(selectedJSONFields, field) + break + } + } + } + if len(selectedJSONFields) > 0 { + conv.parseJSONFields(data[i], selectedJSONFields) + } + } } return &AssistantResponse{ diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go index da3b4892..5b74c8a5 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -775,6 +775,46 @@ func TestXunAssistantPagination(t *testing.T) { assert.Nil(t, err) assert.Greater(t, len(resp.Data), 0) + // Test filtering by assistant_id + // First get an assistant_id from previous results + firstAssistantID := resp.Data[0]["assistant_id"].(string) + + // Test exact match with assistant_id + resp, err = conv.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 = conv.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"]) + // 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 = conv.GetAssistants(AssistantFilter{ + AssistantID: "non-existent-id", + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Equal(t, 0, len(resp.Data)) + // Test combined filters resp, err = conv.GetAssistants(AssistantFilter{ Tags: []string{"tag0"}, @@ -786,4 +826,51 @@ func TestXunAssistantPagination(t *testing.T) { PageSize: 10, }) assert.Nil(t, err) + + // Test filtering with select fields + resp, err = conv.GetAssistants(AssistantFilter{ + Select: []string{"name", "description", "tags"}, + Page: 1, + PageSize: 10, + }) + 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 = conv.GetAssistants(AssistantFilter{ + Tags: []string{"tag0"}, + 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") + } } diff --git a/neo/process.go b/neo/process.go index ac631428..7a27cbac 100644 --- a/neo/process.go +++ b/neo/process.go @@ -22,10 +22,11 @@ func GetNeo() *DSL { func init() { process.RegisterGroup("neo", map[string]process.Handler{ "write": ProcessWrite, - "assistant.add": processAssistantAdd, + "assistant.create": processAssistantCreate, "assistant.save": processAssistantSave, "assistant.delete": processAssistantDelete, "assistant.search": processAssistantSearch, + "assistant.find": processAssistantFind, }) } @@ -56,8 +57,8 @@ func ProcessWrite(process *process.Process) interface{} { return nil } -// processAssistantAdd process the assistant add request -func processAssistantAdd(process *process.Process) interface{} { +// processAssistantCreate process the assistant create request +func processAssistantCreate(process *process.Process) interface{} { process.ValidateArgNums(1) data := process.ArgsMap(0) @@ -68,7 +69,7 @@ func processAssistantAdd(process *process.Process) interface{} { id, err := neo.Conversation.SaveAssistant(data) if err != nil { - exception.New("Failed to add assistant: %s", 500, err.Error()).Throw() + exception.New("Failed to create assistant: %s", 500, err.Error()).Throw() } return id @@ -175,3 +176,31 @@ func processAssistantSearch(process *process.Process) interface{} { return res } + +// processAssistantFind process the assistant find request +func processAssistantFind(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() + } + + filter := conversation.AssistantFilter{ + AssistantID: assistantID, + Page: 1, + PageSize: 1, + } + + res, err := neo.Conversation.GetAssistants(filter) + if err != nil { + exception.New("Failed to find assistant: %s", 500, err.Error()).Throw() + } + + if len(res.Data) == 0 { + exception.New("Assistant not found: %s", 404, assistantID).Throw() + } + + return res.Data[0] +} diff --git a/neo/process_test.go b/neo/process_test.go index 7dce8f1c..5f1e43e9 100644 --- a/neo/process_test.go +++ b/neo/process_test.go @@ -83,8 +83,8 @@ func TestProcessAssistantCRUD(t *testing.T) { "automated": true, } - // Test processAssistantAdd with string JSON - p, err := process.Of("neo.assistant.add", assistant) + // Test processAssistantCreate with string JSON + p, err := process.Of("neo.assistant.create", assistant) if err != nil { t.Fatal(err) } @@ -97,6 +97,33 @@ func TestProcessAssistantCRUD(t *testing.T) { assistantID := output assert.NotNil(t, assistantID) + // Test processAssistantFind + p, err = process.Of("neo.assistant.find", assistantID) + if err != nil { + t.Fatal(err) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + foundAssistant := output.(map[string]interface{}) + assert.Equal(t, assistantID, foundAssistant["assistant_id"]) + assert.Equal(t, "Test Assistant", foundAssistant["name"]) + assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, foundAssistant["tags"]) + assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, foundAssistant["options"]) + + // Test processAssistantFind with non-existent ID + p, err = process.Of("neo.assistant.find", "non-existent-id") + if err != nil { + t.Fatal(err) + } + + _, err = p.Exec() + assert.NotNil(t, err) + assert.Contains(t, err.Error(), "Assistant not found") + // Test with native type JSON fields assistant2 := map[string]interface{}{ "name": "Test Assistant 2", @@ -115,8 +142,8 @@ func TestProcessAssistantCRUD(t *testing.T) { "automated": true, } - // Test processAssistantAdd with native types - p, err = process.Of("neo.assistant.add", assistant2) + // Test processAssistantCreate with native types + p, err = process.Of("neo.assistant.create", assistant2) if err != nil { t.Fatal(err) } @@ -146,8 +173,8 @@ func TestProcessAssistantCRUD(t *testing.T) { "automated": true, } - // Test processAssistantAdd with nil fields - p, err = process.Of("neo.assistant.add", assistant3) + // Test processAssistantCreate with nil fields + p, err = process.Of("neo.assistant.create", assistant3) if err != nil { t.Fatal(err) } @@ -324,7 +351,7 @@ func TestProcessAssistantSearchPagination(t *testing.T) { "automated": i%3 == 0, } - p, err := process.Of("neo.assistant.add", assistant) + p, err := process.Of("neo.assistant.create", assistant) if err != nil { t.Fatal(err) } @@ -438,7 +465,7 @@ func TestProcessAssistantValidation(t *testing.T) { defer test.Clean() // Test missing required fields - p, err := process.Of("neo.assistant.add", map[string]interface{}{}) + p, err := process.Of("neo.assistant.create", map[string]interface{}{}) if err != nil { t.Fatal(err) } @@ -455,6 +482,16 @@ func TestProcessAssistantValidation(t *testing.T) { _, err = p.Exec() assert.NotNil(t, err) + // Test invalid assistant ID for find + p, err = process.Of("neo.assistant.find", "non-existent-id") + if err != nil { + t.Fatal(err) + } + + _, err = p.Exec() + assert.NotNil(t, err) + assert.Contains(t, err.Error(), "Assistant not found") + // Test invalid page number p, err = process.Of("neo.assistant.search", map[string]interface{}{ "page": -1,