From 90967f58c14b71aa96154799970ca37f5cc2d5a3 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 6 Jan 2025 14:20:39 +0800 Subject: [PATCH] Refactor Neo API context handling and improve context management - Replaced the custom context creation functions with a new context package, enhancing consistency and maintainability across the API. - Updated multiple methods in the DSL to utilize the new context package, ensuring a unified approach to context management. - Removed the obsolete context.go file, streamlining the codebase and reducing redundancy. These changes improve the overall structure and clarity of the Neo API, paving the way for future enhancements in context handling and assistant functionalities. --- neo/api.go | 15 ++++---- neo/assistant/hooks.go | 1 + neo/context.go | 49 ------------------------- neo/context/context.go | 82 ++++++++++++++++++++++++++++++++++++++++++ neo/hooks.go | 7 ++-- neo/neo.go | 27 +++++++------- neo/types.go | 18 ---------- 7 files changed, 109 insertions(+), 90 deletions(-) create mode 100644 neo/assistant/hooks.go delete mode 100644 neo/context.go create mode 100644 neo/context/context.go diff --git a/neo/api.go b/neo/api.go index aebe7bad..a67ad151 100644 --- a/neo/api.go +++ b/neo/api.go @@ -15,6 +15,7 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/process" "github.com/yaoapp/yao/helper" + chatctx "github.com/yaoapp/yao/neo/context" "github.com/yaoapp/yao/neo/message" "github.com/yaoapp/yao/neo/store" ) @@ -173,7 +174,7 @@ func (neo *DSL) handleUpload(c *gin.Context) { } // Set the context - ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "") + ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "") defer cancel() // Upload the file @@ -214,7 +215,7 @@ func (neo *DSL) handleChat(c *gin.Context) { } // Set the context with validated chat_id - ctx, cancel := NewContextWithCancel(sid, chatID, c.Query("context")) + ctx, cancel := chatctx.NewWithCancel(sid, chatID, c.Query("context")) defer cancel() neo.Answer(ctx, content, c) @@ -297,7 +298,7 @@ func (neo *DSL) handleDownload(c *gin.Context) { } // Set the context - ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "") + ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "") defer cancel() // Download the file @@ -537,7 +538,7 @@ func (neo *DSL) handleChatUpdate(c *gin.Context) { // If content is not empty, Generate the chat title if body.Content != "" { - ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "") + ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "") defer cancel() title, err := neo.GenerateChatTitle(ctx, body.Content, c, true) @@ -729,7 +730,7 @@ func (neo *DSL) handleGenerateTitle(c *gin.Context) { return } - ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "") defer cancel() // Use silent mode for regular HTTP requests, streaming for SSE @@ -780,7 +781,7 @@ func (neo *DSL) handleGeneratePrompts(c *gin.Context) { return } - ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "") defer cancel() // Use silent mode for regular HTTP requests, streaming for SSE @@ -831,7 +832,7 @@ func (neo *DSL) handleGenerateCustom(c *gin.Context) { return } - ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "") + ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "") defer cancel() // Use silent mode for regular HTTP requests, streaming for SSE diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go new file mode 100644 index 00000000..31df51cd --- /dev/null +++ b/neo/assistant/hooks.go @@ -0,0 +1 @@ +package assistant diff --git a/neo/context.go b/neo/context.go deleted file mode 100644 index ad848c95..00000000 --- a/neo/context.go +++ /dev/null @@ -1,49 +0,0 @@ -package neo - -import ( - "context" - "time" - - jsoniter "github.com/json-iterator/go" - "github.com/yaoapp/kun/log" -) - -// NewContext create a new context -func NewContext(sid, cid, payload string) Context { - ctx := Context{Context: context.Background(), Sid: sid, ChatID: cid} - if payload == "" { - return ctx - } - - err := jsoniter.Unmarshal([]byte(payload), &ctx) - if err != nil { - log.Error("%s", err.Error()) - } - return ctx -} - -// NewContextWithCancel create a new context with cancel -func NewContextWithCancel(sid, cid, payload string) (Context, context.CancelFunc) { - ctx := NewContext(sid, cid, payload) - return ContextWithCancel(ctx) -} - -// NewContextWithTimeout create a new context with timeout -func NewContextWithTimeout(sid, cid, payload string, timeout time.Duration) (Context, context.CancelFunc) { - ctx := NewContext(sid, cid, payload) - return ContextWithTimeout(ctx, timeout) -} - -// ContextWithCancel create a new context -func ContextWithCancel(parent Context) (Context, context.CancelFunc) { - new, cancel := context.WithCancel(parent.Context) - parent.Context = new - return parent, cancel -} - -// ContextWithTimeout create a new context -func ContextWithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) { - new, cancel := context.WithTimeout(parent.Context, timeout) - parent.Context = new - return parent, cancel -} diff --git a/neo/context/context.go b/neo/context/context.go new file mode 100644 index 00000000..79b9ff4b --- /dev/null +++ b/neo/context/context.go @@ -0,0 +1,82 @@ +package context + +import ( + "context" + "time" + + jsoniter "github.com/json-iterator/go" + "github.com/yaoapp/kun/log" +) + +// Context the context +type Context struct { + context.Context + Sid string `json:"sid" yaml:"-"` // Session ID + ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat + AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant + Stack string `json:"stack,omitempty"` + Path string `json:"pathname,omitempty"` + FormData map[string]interface{} `json:"formdata,omitempty"` + Field *Field `json:"field,omitempty"` + Namespace string `json:"namespace,omitempty"` + Config map[string]interface{} `json:"config,omitempty"` + Signal interface{} `json:"signal,omitempty"` + Upload *FileUpload `json:"upload,omitempty"` +} + +// Field the context field +type Field struct { + Name string `json:"name,omitempty"` + Type string `json:"type,omitempty"` + Bind string `json:"bind,omitempty"` + Props map[string]interface{} `json:"props,omitempty"` + Children []interface{} `json:"children,omitempty"` +} + +// FileUpload the file upload +type FileUpload struct { + Name string `json:"name,omitempty"` + Type string `json:"type,omitempty"` + Size int64 `json:"size,omitempty"` + TempFile string `json:"temp_file,omitempty"` +} + +// New create a new context +func New(sid, cid, payload string) Context { + ctx := Context{Context: context.Background(), Sid: sid, ChatID: cid} + if payload == "" { + return ctx + } + + err := jsoniter.Unmarshal([]byte(payload), &ctx) + if err != nil { + log.Error("%s", err.Error()) + } + return ctx +} + +// NewWithCancel create a new context with cancel +func NewWithCancel(sid, cid, payload string) (Context, context.CancelFunc) { + ctx := New(sid, cid, payload) + return WithCancel(ctx) +} + +// NewWithTimeout create a new context with timeout +func NewWithTimeout(sid, cid, payload string, timeout time.Duration) (Context, context.CancelFunc) { + ctx := New(sid, cid, payload) + return WithTimeout(ctx, timeout) +} + +// WithCancel create a new context +func WithCancel(parent Context) (Context, context.CancelFunc) { + new, cancel := context.WithCancel(parent.Context) + parent.Context = new + return parent, cancel +} + +// WithTimeout create a new context +func WithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) { + new, cancel := context.WithTimeout(parent.Context, timeout) + parent.Context = new + return parent, cancel +} diff --git a/neo/hooks.go b/neo/hooks.go index 4ec7dabb..9a3b6e58 100644 --- a/neo/hooks.go +++ b/neo/hooks.go @@ -7,10 +7,11 @@ import ( "github.com/gin-gonic/gin" jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/process" + chatctx "github.com/yaoapp/yao/neo/context" ) // HookCreate create the assistant -func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) { +func (neo *DSL) HookCreate(ctx chatctx.Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) { // Default assistant assistantID := neo.Use @@ -69,7 +70,7 @@ func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gi } // HookPrepare executes the prepare hook before AI is called -func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]map[string]interface{}, error) { +func (neo *DSL) HookPrepare(ctx chatctx.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) { if neo.Prepare == "" { return messages, nil } @@ -114,7 +115,7 @@ func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]m } // HookWrite executes the write hook when response is received from AI -func (neo *DSL) HookWrite(ctx Context, messages []map[string]interface{}, response map[string]interface{}, content string, writer *gin.ResponseWriter) ([]map[string]interface{}, error) { +func (neo *DSL) HookWrite(ctx chatctx.Context, messages []map[string]interface{}, response map[string]interface{}, content string, writer *gin.ResponseWriter) ([]map[string]interface{}, error) { if neo.Write == "" { return []map[string]interface{}{response}, nil } diff --git a/neo/neo.go b/neo/neo.go index f9d56b35..f374bc1d 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -9,6 +9,7 @@ import ( "github.com/gin-gonic/gin" "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/neo/assistant" + chatctx "github.com/yaoapp/yao/neo/context" "github.com/yaoapp/yao/neo/message" ) @@ -16,7 +17,7 @@ import ( var lock sync.Mutex = sync.Mutex{} // Answer reply the message -func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error { +func (neo *DSL) Answer(ctx chatctx.Context, question string, c *gin.Context) error { messages, err := neo.chatMessages(ctx, question) if err != nil { msg := message.New().Error(err).Done() @@ -50,7 +51,7 @@ func (neo *DSL) Select(id string) (assistant.API, error) { } // GeneratePrompts generate prompts for the AI assistant -func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) { +func (neo *DSL) GeneratePrompts(ctx chatctx.Context, input string, c *gin.Context, silent ...bool) (string, error) { prompts := ` Optimize the prompts for the AI assistant 1. Optimize prompts based on the user's input @@ -68,7 +69,7 @@ func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silen } // GenerateChatTitle generate the chat title -func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) { +func (neo *DSL) GenerateChatTitle(ctx chatctx.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. @@ -84,7 +85,7 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, sil } // 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) { +func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) { messages := []map[string]interface{}{ {"role": "system", "content": systemPrompt}, { @@ -187,7 +188,7 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy } // Upload upload a file -func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) { +func (neo *DSL) Upload(ctx chatctx.Context, c *gin.Context) (*assistant.File, error) { // Get the file tmpfile, err := c.FormFile("file") if err != nil { @@ -212,11 +213,11 @@ func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) { } // Get file info - ctx.Upload = &FileUpload{ - Bytes: int(tmpfile.Size), - Name: tmpfile.Filename, - ContentType: tmpfile.Header.Get("Content-Type"), - Option: option, + ctx.Upload = &chatctx.FileUpload{ + Name: tmpfile.Filename, + Type: tmpfile.Header.Get("Content-Type"), + Size: tmpfile.Size, + TempFile: tmpfile.Filename, } // Default use the assistant in context @@ -235,7 +236,7 @@ func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) { } // Download downloads a file -func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse, error) { +func (neo *DSL) Download(ctx chatctx.Context, c *gin.Context) (*assistant.FileResponse, error) { // Get file_id from query string fileID := c.Query("file_id") if fileID == "" { @@ -259,7 +260,7 @@ func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse, } // chat chat with AI -func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]interface{}, c *gin.Context) error { +func (neo *DSL) chat(ast assistant.API, ctx chatctx.Context, messages []map[string]interface{}, c *gin.Context) error { if ast == nil { msg := message.New().Error("assistant is not initialized").Done() msg.Write(c.Writer) @@ -339,7 +340,7 @@ func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]inter } // chatMessages get the chat messages -func (neo *DSL) chatMessages(ctx Context, content ...string) ([]map[string]interface{}, error) { +func (neo *DSL) chatMessages(ctx chatctx.Context, content ...string) ([]map[string]interface{}, error) { history, err := neo.Store.GetHistory(ctx.Sid, ctx.ChatID) if err != nil { diff --git a/neo/types.go b/neo/types.go index e3511168..bbf614d2 100644 --- a/neo/types.go +++ b/neo/types.go @@ -1,8 +1,6 @@ package neo import ( - "context" - "github.com/gin-gonic/gin" "github.com/yaoapp/yao/neo/assistant" "github.com/yaoapp/yao/neo/rag" @@ -48,22 +46,6 @@ type Mention struct { Type string `json:"type,omitempty"` } -// Context the context -type Context struct { - Sid string `json:"sid" yaml:"-"` // Session ID - ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat - AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant - Stack string `json:"stack,omitempty"` - Path string `json:"pathname,omitempty"` - FormData map[string]interface{} `json:"formdata,omitempty"` - Field *Field `json:"field,omitempty"` - Namespace string `json:"namespace,omitempty"` - Config map[string]interface{} `json:"config,omitempty"` - Signal interface{} `json:"signal,omitempty"` - Upload *FileUpload `json:"upload,omitempty"` - context.Context `json:"-" yaml:"-"` -} - // Field the context field type Field struct { Name string `json:"name,omitempty"`