diff --git a/neo/api.go b/neo/api.go index 2fd9289d..8f692db6 100644 --- a/neo/api.go +++ b/neo/api.go @@ -61,6 +61,11 @@ func (neo *DSL) API(router *gin.Engine, path string) error { // Mention api router.GET(path+"/mentions", append(middlewares, neo.handleMentions)...) + // Generate api + router.POST(path+"/generate", append(middlewares, neo.handleGenerateCustom)...) + router.POST(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...) + router.POST(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) + // Dangerous operations router.DELETE(path+"/dangerous/clear_chats", append(middlewares, neo.handleChatsDeleteAll)...) @@ -390,29 +395,8 @@ func (neo *DSL) handleMentions(c *gin.Context) { return } - // Add test data - testMentions := []Mention{ - { - ID: "assistant_1", - Name: "Alice AI", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Alice", - }, - { - ID: "assistant_2", - Name: "Bob Bot", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Bob", - }, - { - ID: "assistant_3", - Name: "Carol AI", - Type: "assistant", - Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Carol", - }, - } - // Filter mentions by keywords + testMentions := TestMentions if keywords != "" { filtered := []Mention{} for _, m := range testMentions { @@ -534,3 +518,142 @@ func (neo *DSL) handleChatsDeleteAll(c *gin.Context) { c.JSON(200, gin.H{"message": "ok"}) c.Done() } + +// generateResponse is a helper struct to handle both SSE and HTTP responses +type generateResponse struct { + c *gin.Context + sid string + content string + result interface{} + err error +} + +// validate checks common validation rules +func (r *generateResponse) validate() bool { + if r.sid == "" { + r.c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + return false + } + + if r.content == "" { + r.c.JSON(400, gin.H{"message": "content is required", "code": 400}) + return false + } + + return true +} + +// send handles both SSE and HTTP responses +func (r *generateResponse) send(key string) { + if r.err != nil { + if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { + msg := message.New().Error(r.err.Error()).Done() + msg.Write(r.c.Writer) + } else { + r.c.JSON(500, gin.H{"message": r.err.Error(), "code": 500}) + } + return + } + + if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { + // Set headers for SSE + r.c.Header("Content-Type", "text/event-stream;charset=utf-8") + r.c.Header("Cache-Control", "no-cache") + r.c.Header("Connection", "keep-alive") + + msg := message.New().Map(gin.H{key: r.result}).Done() + msg.Write(r.c.Writer) + } else { + r.c.JSON(200, gin.H{key: r.result}) + } +} + +// handleGenerateTitle handles generating a chat title +func (neo *DSL) handleGenerateTitle(c *gin.Context) { + var body struct { + Content string `json:"content"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + + resp := &generateResponse{ + c: c, + sid: c.GetString("__sid"), + content: body.Content, + } + if !resp.validate() { + return + } + + ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + defer cancel() + + resp.result, resp.err = neo.GenerateChatTitle(ctx, resp.content, c) + resp.send("result") +} + +// handleGeneratePrompts handles generating prompts +func (neo *DSL) handleGeneratePrompts(c *gin.Context) { + var body struct { + Content string `json:"content"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + + resp := &generateResponse{ + c: c, + sid: c.GetString("__sid"), + content: body.Content, + } + if !resp.validate() { + return + } + + ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + defer cancel() + + resp.result, resp.err = neo.GeneratePrompts(ctx, resp.content, c) + resp.send("result") +} + +// handleGenerateCustom handles generating custom content +func (neo *DSL) handleGenerateCustom(c *gin.Context) { + var body struct { + Content string `json:"content"` + Type string `json:"type"` + SystemPrompt string `json:"system_prompt"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + + resp := &generateResponse{ + c: c, + sid: c.GetString("__sid"), + content: body.Content, + } + if !resp.validate() { + return + } + + // Additional validations for custom generation + if body.Type == "" { + c.JSON(400, gin.H{"message": "type is required", "code": 400}) + return + } + if body.SystemPrompt == "" { + c.JSON(400, gin.H{"message": "system_prompt is required", "code": 400}) + return + } + + ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + defer cancel() + + resp.result, resp.err = neo.GenerateWithAI(ctx, resp.content, body.Type, body.SystemPrompt, c) + resp.send("result") +} diff --git a/neo/neo.go b/neo/neo.go index 48ef866e..115ebc97 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -58,20 +58,40 @@ func (neo *DSL) GetMentions(keywords string) ([]Mention, error) { return neo.HookMention(context.Background(), keywords) } +// GeneratePrompts generate prompts for the AI assistant +func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context) (string, error) { + prompts := ` + Help me generate prompts for the AI assistant + 1. The prompts should guide the AI to better understand and respond to user questions + 2. The prompts should be clear and specific + 3. The prompts should be in the same language as the input + 4. Keep the prompts concise but comprehensive + ` + return neo.GenerateWithAI(ctx, input, "prompts", prompts, c) +} + // GenerateChatTitle generate the chat title func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (string, error) { - prompts := ` - Help me generate a title for the chat - 1. The title should be a short and concise description of the chat. - 2. The title should be a single sentence. - 3. The title should be in same language as the chat. - 4. The title should be no more than 50 characters. + Help me generate a title for the chat + 1. The title should be a short and concise description of the chat. + 2. The title should be a single sentence. + 3. The title should be in same language as the chat. + 4. The title should be no more than 50 characters. ` + return neo.GenerateWithAI(ctx, input, "title", prompts, c) +} +// GenerateWithAI generate content with AI, type can be "title", "prompts", etc. +func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context) (string, error) { messages := []map[string]interface{}{ - {"role": "system", "content": prompts}, - {"role": "user", "content": input}, + {"role": "system", "content": systemPrompt}, + { + "role": "user", + "content": input, + "type": messageType, + "name": ctx.Sid, + }, } res, err := neo.HookCreate(ctx, messages, c) diff --git a/neo/test_data.go b/neo/test_data.go new file mode 100644 index 00000000..f783b8de --- /dev/null +++ b/neo/test_data.go @@ -0,0 +1,30 @@ +package neo + +// TestMentions contains test data for mentions +var TestMentions = []Mention{ + { + ID: "assistant_1", + Name: "Alice AI", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Alice", + }, + { + ID: "assistant_2", + Name: "Bob Bot", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Bob", + }, + { + ID: "assistant_3", + Name: "Carol AI", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Carol", + }, + // ... 其他测试数据 ... + { + ID: "assistant_20", + Name: "Tara Tech", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Tara", + }, +}