diff --git a/neo/api.go b/neo/api.go index 2fd9289d..be5640a4 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" @@ -36,32 +37,120 @@ 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) + 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)...) - // Chat 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 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)...) + // 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)...) + + // 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' router.DELETE(path+"/dangerous/clear_chats", append(middlewares, neo.handleChatsDeleteAll)...) return nil @@ -383,49 +472,35 @@ 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 } - // 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 - 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.Data { + mention := Mention{ + ID: item["assistant_id"].(string), + Name: item["name"].(string), + Type: item["type"].(string), + Avatar: item["avatar"].(string), } - testMentions = filtered + mentions = append(mentions, mention) } - // Append test data to actual mentions - mentions = append(mentions, testMentions...) - c.JSON(200, map[string]interface{}{"data": mentions}) c.Done() } @@ -462,7 +537,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() @@ -534,3 +609,417 @@ 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 == "" { + 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 == "" { + 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 + } + + 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") { + 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}) + } + return + } + + 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(). + 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 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: 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 + } + + ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + defer cancel() + + // 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 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: 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 + } + + ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + defer cancel() + + // 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 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: content, + } + if !resp.validate() { + return + } + + // Additional validations for custom generation + if genType == "" { + c.JSON(400, gin.H{"message": "type is required", "code": 400}) + return + } + if 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() + + // 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") +} + +// 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, ",") + } + + // 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}) + c.Done() + return + } + + c.JSON(200, response) + 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") + if assistantID == "" { + c.JSON(400, gin.H{"message": "assistant id is required", "code": 400}) + c.Done() + return + } + + filter := conversation.AssistantFilter{ + AssistantID: assistantID, + 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 + } + + 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": response.Data[0]}) + 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 + } + + 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() +} + +// 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() +} + +// 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() +} 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/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..d4b20a16 100644 --- a/neo/conversation/mongo.go +++ b/neo/conversation/mongo.go @@ -1,60 +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) 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 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 +} + +// 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 afb1aa74..fb150570 100644 --- a/neo/conversation/redis.go +++ b/neo/conversation/redis.go @@ -1,60 +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) 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 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 +} + +// 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 205a99d6..521fc176 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -1,52 +1,133 @@ package conversation -// 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 } -// Conversation the store interface -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 - DeleteChat(sid string, cid string) error - DeleteAllChats(sid string) error - UpdateChatTitle(sid string, cid string, title string) error +// AssistantFilter represents the assistant filter structure +// Used for filtering and pagination when retrieving assistant lists +type AssistantFilter struct { + Tags []string `json:"tags,omitempty"` // Filter by tags + Keywords string `json:"keywords,omitempty"` // Search in name and description + Connector string `json:"connector,omitempty"` // Filter by connector + AssistantID string `json:"assistant_id,omitempty"` // Filter by assistant ID + Mentionable *bool `json:"mentionable,omitempty"` // Filter by mentionable status + Automated *bool `json:"automated,omitempty"` // Filter by automation status + 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 +// Used for returning paginated assistant lists +type AssistantResponse struct { + 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 +// 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{}) (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/weaviate.go b/neo/conversation/weaviate.go index 7b251851..8462e843 100644 --- a/neo/conversation/weaviate.go +++ b/neo/conversation/weaviate.go @@ -1,60 +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) 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 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 +} + +// 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 61a25b55..d6ee4efa 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" @@ -15,32 +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 } -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"` -} - -// 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) { @@ -111,6 +114,11 @@ func (conv *Xun) initialize() error { return err } + // Initialize assistant table + if err := conv.initAssistantTable(); err != nil { + return err + } + return nil } @@ -126,11 +134,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 +157,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 +207,59 @@ 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() // 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.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 + 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() + }) + + 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", "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) + } + } + + return nil +} + func (conv *Xun) getUserID(sid string) (string, error) { field := "user_id" if conv.setting.UserField != "" { @@ -217,13 +279,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 +446,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 +467,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 +521,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 +575,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) @@ -638,3 +661,265 @@ func (conv *Xun) DeleteAllChats(sid string) error { Delete() 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{}) (interface{}, error) { + // Validate required fields + requiredFields := []string{"name", "type", "connector"} + for _, field := range requiredFields { + if _, ok := assistant[field]; !ok { + return nil, fmt.Errorf("field %s is required", field) + } + if assistant[field] == nil || assistant[field] == "" { + return nil, fmt.Errorf("field %s cannot be empty", field) + } + } + + // 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 := 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", assistantCopy["assistant_id"]). + Exists() + if err != nil { + 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 { + _, err := conv.query.New(). + Table(conv.getAssistantTable()). + Where("assistant_id", assistantCopy["assistant_id"]). + Update(assistantCopy) + if err != nil { + return nil, err + } + return assistantCopy["assistant_id"], nil + } + + 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 +func (conv *Xun) DeleteAssistant(assistantID string) error { + // 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() + return err +} + +// 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 { + 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", pattern) + } else { + qb.OrWhere("tags", "like", pattern) + } + } + }) + } + + // Apply keyword filter if provided + if filter.Keywords != "" { + qb.Where(func(qb query.Query) { + qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Keywords)). + OrWhere("description", "like", fmt.Sprintf("%%%s%%", filter.Keywords)) + }) + } + + // Apply connector filter if provided + if filter.Connector != "" { + qb.Where("connector", filter.Connector) + } + + // Apply assistant_id filter if provided + if filter.AssistantID != "" { + qb.Where("assistant_id", filter.AssistantID) + } + + // Apply mentionable filter if provided + if filter.Mentionable != nil { + qb.Where("mentionable", *filter.Mentionable) + } + + // Apply automated filter if provided + if filter.Automated != nil { + qb.Where("automated", *filter.Automated) + } + + // Set defaults for pagination + if filter.PageSize <= 0 { + filter.PageSize = 20 + } + if filter.Page <= 0 { + filter.Page = 1 + } + + // Get total count + total, err := qb.Clone().Count() + if err != nil { + return nil, err + } + + // 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 + } + + // 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). + Limit(filter.PageSize). + Get() + if err != nil { + return nil, err + } + + // 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 + // 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{ + 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 d42f32e2..5b74c8a5 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -5,6 +5,7 @@ import ( "testing" "time" + jsoniter "github.com/json-iterator/go" "github.com/stretchr/testify/assert" "github.com/yaoapp/gou/connector" "github.com/yaoapp/xun/capsule" @@ -15,13 +16,28 @@ 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) + } + + // Add a small delay to ensure table is created + time.Sleep(100 * time.Millisecond) + conv, err := NewXun(Setting{ Connector: "default", Table: "__unit_test_conversation", @@ -32,35 +48,59 @@ 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) } - 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)) } - 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 +117,17 @@ 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") + + // Add a small delay to ensure table is created + time.Sleep(100 * time.Millisecond) - sch.DropTableIfExists("__unit_test_conversation") conv, err := NewXun(Setting{ Connector: "mysql", Table: "__unit_test_conversation", @@ -90,44 +138,73 @@ 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) } - 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)) } - 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) } @@ -143,7 +220,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,44 +231,18 @@ 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() - 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) } @@ -209,7 +260,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 +275,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 @@ -252,11 +303,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) } @@ -294,7 +345,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) } @@ -325,7 +376,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{ @@ -344,7 +395,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 @@ -365,7 +416,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{ @@ -385,7 +436,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) } @@ -403,3 +454,423 @@ 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) + } + + // Add a small delay to ensure table is created + time.Sleep(100 * time.Millisecond) + + conv, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + if err != nil { + t.Fatal(err) + } + + // 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", + "avatar": "https://example.com/avatar.png", + "connector": "openai", + "description": "Test Description", + "tags": tagsJSON, + "options": optionsJSON, + "mentionable": true, + "automated": true, + } + + // Test SaveAssistant (Create) with string JSON + v, err := conv.SaveAssistant(assistant) + assert.Nil(t, err) + assistantID := v.(string) + assert.NotEmpty(t, assistantID) + + // 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, 3, len(resp.Data)) + + // Verify first assistant (string JSON) + found := false + for _, item := range resp.Data { + if item["assistant_id"].(string) == assistantID { + found = true + // Now we expect parsed JSON values instead of JSON strings + assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"]) + assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"]) + break + } + } + assert.True(t, found) + + // Verify second assistant (native types converted to JSON) + found = false + for _, item := range resp.Data { + if item["assistant_id"].(string) == assistant2ID { + found = true + // Now we expect parsed JSON values directly + assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"]) + assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"]) + + // Verify other JSON fields + assert.Equal(t, []interface{}{"prompt1", "prompt2"}, item["prompts"]) + assert.Equal(t, []interface{}{"flow1", "flow2"}, item["flows"]) + assert.Equal(t, []interface{}{"file1", "file2"}, item["files"]) + assert.Equal(t, + []interface{}{ + map[string]interface{}{"name": "func1"}, + map[string]interface{}{"name": "func2"}, + }, + item["functions"]) + assert.Equal(t, + map[string]interface{}{ + "read": true, + "write": true, + }, + item["permissions"]) + break + } + } + assert.True(t, found) + + // Verify third assistant (nil fields) + found = false + for _, item := range resp.Data { + if item["assistant_id"].(string) == assistant3ID { + found = true + assert.Nil(t, item["tags"]) + assert.Nil(t, item["options"]) + assert.Nil(t, item["prompts"]) + assert.Nil(t, item["flows"]) + assert.Nil(t, item["files"]) + assert.Nil(t, item["functions"]) + assert.Nil(t, item["permissions"]) + break + } + } + assert.True(t, found) + + // 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) + 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) + 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 + err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") + if err != nil { + 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", + }) + if err != nil { + t.Fatal(err) + } + + // 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": 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) + } + + // Test first page + resp, err := conv.GetAssistants(AssistantFilter{ + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Equal(t, 10, len(resp.Data)) + assert.Equal(t, int64(25), resp.Total) + assert.Equal(t, 3, resp.PageCnt) + assert.Equal(t, 2, resp.Next) + assert.Equal(t, 0, resp.Prev) + + // Test second page + resp, err = conv.GetAssistants(AssistantFilter{ + Page: 2, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Equal(t, 10, len(resp.Data)) + assert.Equal(t, 3, resp.Next) + assert.Equal(t, 1, resp.Prev) + + // Test last page + resp, err = conv.GetAssistants(AssistantFilter{ + Page: 3, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Equal(t, 5, len(resp.Data)) + assert.Equal(t, 0, resp.Next) + assert.Equal(t, 2, resp.Prev) + + // Test filtering with tags + resp, err = conv.GetAssistants(AssistantFilter{ + Tags: []string{"tag0"}, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Equal(t, 5, len(resp.Data)) + + // Test filtering with keywords + resp, err = conv.GetAssistants(AssistantFilter{ + Keywords: "Assistant 1", + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Greater(t, len(resp.Data), 0) + + // Test filtering with connector + resp, err = conv.GetAssistants(AssistantFilter{ + Connector: "connector0", + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Greater(t, len(resp.Data), 0) + + // Test filtering with mentionable + mentionableTrue := true + resp, err = conv.GetAssistants(AssistantFilter{ + Mentionable: &mentionableTrue, + Page: 1, + PageSize: 10, + }) + assert.Nil(t, err) + assert.Greater(t, len(resp.Data), 0) + + // Test filtering with automated + automatedTrue := true + resp, err = conv.GetAssistants(AssistantFilter{ + Automated: &automatedTrue, + Page: 1, + PageSize: 10, + }) + 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"}, + Keywords: "Assistant", + Connector: "connector0", + Mentionable: &mentionableTrue, + Automated: &automatedTrue, + Page: 1, + 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/neo.go b/neo/neo.go index 48ef866e..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" @@ -58,20 +58,50 @@ func (neo *DSL) GetMentions(keywords string) ([]Mention, error) { return neo.HookMention(context.Background(), keywords) } -// GenerateChatTitle generate the chat title -func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (string, error) { - +// GeneratePrompts generate prompts for the AI assistant +func (neo *DSL) GeneratePrompts(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. - 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. + 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 ` + 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, 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. + 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. + ` + 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, silent bool) (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) @@ -119,8 +149,21 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (st // 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 } @@ -131,7 +174,9 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (st 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 @@ -372,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 } @@ -453,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 { diff --git a/neo/process.go b/neo/process.go index 598b40ad..7a27cbac 100644 --- a/neo/process.go +++ b/neo/process.go @@ -1,15 +1,32 @@ 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.create": processAssistantCreate, + "assistant.save": processAssistantSave, + "assistant.delete": processAssistantDelete, + "assistant.search": processAssistantSearch, + "assistant.find": processAssistantFind, }) } @@ -39,3 +56,151 @@ func ProcessWrite(process *process.Process) interface{} { return nil } + +// processAssistantCreate process the assistant create request +func processAssistantCreate(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() + } + + id, err := neo.Conversation.SaveAssistant(data) + if err != nil { + exception.New("Failed to create assistant: %s", 500, err.Error()).Throw() + } + + return id +} + +// 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() + } + + id, err := neo.Conversation.SaveAssistant(data) + if err != nil { + exception.New("Failed to save assistant: %s", 500, err.Error()).Throw() + } + + return id +} + +// 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 +} + +// 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 new file mode 100644 index 00000000..5f1e43e9 --- /dev/null +++ b/neo/process_test.go @@ -0,0 +1,513 @@ +package neo + +import ( + "fmt" + "testing" + + "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 with string JSON fields + tagsJSON := `["tag1", "tag2", "tag3"]` + optionsJSON := `{"model": "gpt-4"}` + 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 processAssistantCreate with string JSON + p, err := process.Of("neo.assistant.create", assistant) + if err != nil { + t.Fatal(err) + } + + output, err := p.Exec() + if err != nil { + t.Fatal(err) + } + + 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", + "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 processAssistantCreate with native types + p, err = process.Of("neo.assistant.create", 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 processAssistantCreate with nil fields + p, err = process.Of("neo.assistant.create", 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) + } + + 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(3), total) + + items := searchRes.Get("data") + if items == nil { + items = []map[string]interface{}{} + } + assert.Equal(t, 3, len(items.([]map[string]interface{}))) + + // 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) + } + + output, err = p.Exec() + if err != nil { + t.Fatal(err) + } + + searchRes = any.Of(output).Map() + items = searchRes.Get("data") + 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.True(t, found) + + // 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")) + + // 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) + } + + 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) +} + +func TestProcessAssistantSearchPagination(t *testing.T) { + prepare(t) + defer test.Clean() + + // Create multiple assistants for pagination testing + for i := 0; i < 25; i++ { + 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": []string{fmt.Sprintf("tag%d", i%5)}, + "mentionable": i%2 == 0, + "automated": i%3 == 0, + } + + p, err := process.Of("neo.assistant.create", 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.create", 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 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, + "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)