Enhance Neo API with new generation endpoints and improved SSE handling
- Added new GET and POST endpoints for generating custom content, chat titles, and prompts, expanding the API's capabilities. - Implemented support for Server-Sent Events (SSE) in response handling, allowing for real-time updates and improved user experience. - Refactored validation and error handling to differentiate between regular and SSE requests, ensuring appropriate responses based on the request type. - Updated the GenerateWithAI method to support silent mode for regular HTTP requests, enhancing flexibility in content generation. - Improved overall structure and maintainability of the API by consolidating request handling logic.
This commit is contained in:
Max
parent
fc1cbb2457
commit
a90d1d5a0a
2 changed files with 141 additions and 42 deletions
142
neo/api.go
142
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")
|
||||
}
|
||||
|
|
|
|||
41
neo/neo.go
41
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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue