diff --git a/neo/api.go b/neo/api.go index 9d35b0c0..b680a711 100644 --- a/neo/api.go +++ b/neo/api.go @@ -65,8 +65,11 @@ func (neo *DSL) API(router *gin.Engine, path string) error { router.GET(path+"/mentions", append(middlewares, neo.handleMentions)...) // Generate api + router.GET(path+"/generate", append(middlewares, neo.handleGenerateCustom)...) router.POST(path+"/generate", append(middlewares, neo.handleGenerateCustom)...) + router.GET(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...) router.POST(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...) + router.GET(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) router.POST(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...) // Dangerous operations @@ -534,12 +537,32 @@ type generateResponse struct { // validate checks common validation rules func (r *generateResponse) validate() bool { if r.sid == "" { - r.c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { + r.c.Header("Content-Type", "text/event-stream;charset=utf-8") + r.c.Header("Cache-Control", "no-cache") + r.c.Header("Connection", "keep-alive") + msg := message.New(). + Error("sid is required"). + Done() + msg.Write(r.c.Writer) + } else { + r.c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + } return false } if r.content == "" { - r.c.JSON(400, gin.H{"message": "content is required", "code": 400}) + if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { + r.c.Header("Content-Type", "text/event-stream;charset=utf-8") + r.c.Header("Cache-Control", "no-cache") + r.c.Header("Connection", "keep-alive") + msg := message.New(). + Error("content is required"). + Done() + msg.Write(r.c.Writer) + } else { + r.c.JSON(400, gin.H{"message": "content is required", "code": 400}) + } return false } @@ -550,7 +573,12 @@ func (r *generateResponse) validate() bool { func (r *generateResponse) send(key string) { if r.err != nil { if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { - msg := message.New().Error(r.err.Error()).Done() + r.c.Header("Content-Type", "text/event-stream;charset=utf-8") + r.c.Header("Cache-Control", "no-cache") + r.c.Header("Connection", "keep-alive") + msg := message.New(). + Error(r.err.Error()). + Done() msg.Write(r.c.Writer) } else { r.c.JSON(500, gin.H{"message": r.err.Error(), "code": 500}) @@ -559,12 +587,12 @@ func (r *generateResponse) send(key string) { } if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") { - // Set headers for SSE r.c.Header("Content-Type", "text/event-stream;charset=utf-8") r.c.Header("Cache-Control", "no-cache") r.c.Header("Connection", "keep-alive") - - msg := message.New().Map(gin.H{key: r.result}).Done() + msg := message.New(). + Map(gin.H{key: r.result}). + Done() msg.Write(r.c.Writer) } else { r.c.JSON(200, gin.H{key: r.result}) @@ -573,18 +601,24 @@ func (r *generateResponse) send(key string) { // handleGenerateTitle handles generating a chat title func (neo *DSL) handleGenerateTitle(c *gin.Context) { - var body struct { - Content string `json:"content"` - } - if err := c.BindJSON(&body); err != nil { - c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) - return + var content string + if c.Request.Method == "GET" { + content = c.Query("content") + } else { + var body struct { + Content string `json:"content"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + content = body.Content } resp := &generateResponse{ c: c, sid: c.GetString("__sid"), - content: body.Content, + content: content, } if !resp.validate() { return @@ -593,25 +627,50 @@ func (neo *DSL) handleGenerateTitle(c *gin.Context) { ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") defer cancel() - resp.result, resp.err = neo.GenerateChatTitle(ctx, resp.content, c) + // Use silent mode for regular HTTP requests, streaming for SSE + silent := !strings.Contains(c.GetHeader("Accept"), "text/event-stream") + resp.result, resp.err = neo.GenerateChatTitle(ctx, resp.content, c, silent) resp.send("result") } // handleGeneratePrompts handles generating prompts func (neo *DSL) handleGeneratePrompts(c *gin.Context) { - var body struct { - Content string `json:"content"` - } - if err := c.BindJSON(&body); err != nil { - c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) - return + var content string + if c.Request.Method == "GET" { + content = c.Query("content") + } else { + var body struct { + Content string `json:"content"` + } + if err := c.BindJSON(&body); err != nil { + // For SSE requests, send error message in SSE format + if strings.Contains(c.GetHeader("Accept"), "text/event-stream") { + c.Header("Content-Type", "text/event-stream;charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + msg := message.New().Error("invalid request body").Done() + msg.Write(c.Writer) + return + } + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + content = body.Content } resp := &generateResponse{ c: c, sid: c.GetString("__sid"), - content: body.Content, + content: content, } + + // For SSE requests, set headers before validation + if strings.Contains(c.GetHeader("Accept"), "text/event-stream") { + c.Header("Content-Type", "text/event-stream;charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + } + if !resp.validate() { return } @@ -619,37 +678,50 @@ func (neo *DSL) handleGeneratePrompts(c *gin.Context) { ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") defer cancel() - resp.result, resp.err = neo.GeneratePrompts(ctx, resp.content, c) + // Use silent mode for regular HTTP requests, streaming for SSE + silent := !strings.Contains(c.GetHeader("Accept"), "text/event-stream") + resp.result, resp.err = neo.GeneratePrompts(ctx, resp.content, c, silent) resp.send("result") } // handleGenerateCustom handles generating custom content func (neo *DSL) handleGenerateCustom(c *gin.Context) { - var body struct { - Content string `json:"content"` - Type string `json:"type"` - SystemPrompt string `json:"system_prompt"` - } - if err := c.BindJSON(&body); err != nil { - c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) - return + var content, genType, systemPrompt string + + if c.Request.Method == "GET" { + content = c.Query("content") + genType = c.Query("type") + systemPrompt = c.Query("system_prompt") + } else { + var body struct { + Content string `json:"content"` + Type string `json:"type"` + SystemPrompt string `json:"system_prompt"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + return + } + content = body.Content + genType = body.Type + systemPrompt = body.SystemPrompt } resp := &generateResponse{ c: c, sid: c.GetString("__sid"), - content: body.Content, + content: content, } if !resp.validate() { return } // Additional validations for custom generation - if body.Type == "" { + if genType == "" { c.JSON(400, gin.H{"message": "type is required", "code": 400}) return } - if body.SystemPrompt == "" { + if systemPrompt == "" { c.JSON(400, gin.H{"message": "system_prompt is required", "code": 400}) return } @@ -657,6 +729,8 @@ func (neo *DSL) handleGenerateCustom(c *gin.Context) { ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") defer cancel() - resp.result, resp.err = neo.GenerateWithAI(ctx, resp.content, body.Type, body.SystemPrompt, c) + // Use silent mode for regular HTTP requests, streaming for SSE + silent := !strings.Contains(c.GetHeader("Accept"), "text/event-stream") + resp.result, resp.err = neo.GenerateWithAI(ctx, resp.content, genType, systemPrompt, c, silent) resp.send("result") } diff --git a/neo/neo.go b/neo/neo.go index 115ebc97..a83ec685 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -59,19 +59,25 @@ func (neo *DSL) GetMentions(keywords string) ([]Mention, error) { } // GeneratePrompts generate prompts for the AI assistant -func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context) (string, error) { +func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) { prompts := ` - Help me generate prompts for the AI assistant - 1. The prompts should guide the AI to better understand and respond to user questions + Optimize the prompts for the AI assistant + 1. Optimize prompts based on the user's input 2. The prompts should be clear and specific 3. The prompts should be in the same language as the input 4. Keep the prompts concise but comprehensive + 5. DO NOT ASK USER FOR MORE INFORMATION, JUST GENERATE PROMPTS + 6. DO NOT ANSWER THE QUESTION, JUST GENERATE PROMPTS ` - return neo.GenerateWithAI(ctx, input, "prompts", prompts, c) + isSilent := false + if len(silent) > 0 { + isSilent = silent[0] + } + return neo.GenerateWithAI(ctx, input, "prompts", prompts, c, isSilent) } // GenerateChatTitle generate the chat title -func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (string, error) { +func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) { prompts := ` Help me generate a title for the chat 1. The title should be a short and concise description of the chat. @@ -79,11 +85,15 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (st 3. The title should be in same language as the chat. 4. The title should be no more than 50 characters. ` - return neo.GenerateWithAI(ctx, input, "title", prompts, c) + isSilent := false + if len(silent) > 0 { + isSilent = silent[0] + } + return neo.GenerateWithAI(ctx, input, "title", prompts, c, isSilent) } // GenerateWithAI generate content with AI, type can be "title", "prompts", etc. -func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context) (string, error) { +func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) { messages := []map[string]interface{}{ {"role": "system", "content": systemPrompt}, { @@ -139,8 +149,21 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy // Append content and send message content = msg.Append(content) + // Only send real-time messages if not in silent mode + if !silent && msg.Message != nil && msg.Message.Text != "" { + message.New(). + Map(map[string]interface{}{ + "text": msg.Message.Text, + "done": msg.Message.Done, + }). + Write(c.Writer) + } + // Complete the stream if msg.Message.Done { + if !silent && msg.Message.Text == "" { + msg.Write(c.Writer) + } done <- true return 0 // break } @@ -151,7 +174,9 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy if err != nil { log.Error("Chat error: %s", err.Error()) - message.New().Error(err).Done().Write(c.Writer) + if !silent { + message.New().Error(err).Done().Write(c.Writer) + } } done <- true