From b01f27d70f0b75c66eb64764ae3a60c3fc50becf Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 13 Jan 2025 11:07:12 +0800 Subject: [PATCH] Refactor assistant initialization and enhance context handling in Neo API - Updated the Answer method to improve assistant initialization by directly calling the new HookInit method, streamlining the process of selecting and initializing assistants based on context. - Introduced ResHookInit struct to encapsulate the response from the assistant initialization hook, enhancing clarity and maintainability. - Enhanced the context package by adding a Map method to facilitate easier mapping of context data, improving the overall structure of context management. - Refactored the assistant interface to include the new HookInit method, ensuring a consistent approach to assistant interactions. These changes improve the robustness and maintainability of the Neo API, paving the way for future enhancements in assistant functionalities and context management. --- neo/assistant/hooks.go | 73 ++++++++++++++++++++++++++++++++++++++++++ neo/assistant/types.go | 3 ++ neo/context/context.go | 40 +++++++++++++++++++++++ neo/neo.go | 27 +++++++++++----- 4 files changed, 135 insertions(+), 8 deletions(-) diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index 31df51cd..b9c8fb7b 100644 --- a/neo/assistant/hooks.go +++ b/neo/assistant/hooks.go @@ -1 +1,74 @@ package assistant + +import ( + "fmt" + + chatctx "github.com/yaoapp/yao/neo/context" + "github.com/yaoapp/yao/neo/message" +) + +const ( + // HookErrorMethodNotFound is the error message for method not found + HookErrorMethodNotFound = "method not found" +) + +// ResHookInit the response of the init hook +type ResHookInit struct { + AssistantID string `json:"assistant_id,omitempty"` + ChatID string `json:"chat_id,omitempty"` +} + +// HookInit initialize the assistant +func (ast *Assistant) HookInit(context chatctx.Context, messages []message.Message) (*ResHookInit, error) { + v, err := ast.call("Init", context, messages) + if err != nil { + if err.Error() == HookErrorMethodNotFound { + return nil, nil + } + return nil, err + } + + response := &ResHookInit{} + switch v := v.(type) { + case map[string]interface{}: + if res, ok := v["assistant_id"].(string); ok { + response.AssistantID = res + } + if res, ok := v["chat_id"].(string); ok { + response.ChatID = res + } + + case string: + response.AssistantID = v + response.ChatID = context.ChatID + + case nil: + response.AssistantID = ast.ID + response.ChatID = context.ChatID + } + + return response, nil +} + +// Call the script method +func (ast *Assistant) call(method string, context chatctx.Context, args ...any) (interface{}, error) { + + if ast.Script == nil { + return nil, nil + } + + ctx, err := ast.Script.NewContext(context.Sid, nil) + if err != nil { + return nil, err + } + defer ctx.Close() + + // Check if the method exists + if !ctx.Global().Has(method) { + return nil, fmt.Errorf(HookErrorMethodNotFound) + } + + // Call the method + args = append([]interface{}{context.Map()}, args...) + return ctx.Call(method, args...) +} diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 2002ae24..71dddf19 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -7,6 +7,8 @@ import ( "github.com/yaoapp/gou/rag/driver" v8 "github.com/yaoapp/gou/runtime/v8" + chatctx "github.com/yaoapp/yao/neo/context" + "github.com/yaoapp/yao/neo/message" api "github.com/yaoapp/yao/openai" ) @@ -16,6 +18,7 @@ type API interface { Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error) Download(ctx context.Context, fileID string) (*FileResponse, error) ReadBase64(ctx context.Context, fileID string) (string, error) + HookInit(ctx chatctx.Context, messages []message.Message) (*ResHookInit, error) } // RAG the RAG interface diff --git a/neo/context/context.go b/neo/context/context.go index 79b9ff4b..da0efd51 100644 --- a/neo/context/context.go +++ b/neo/context/context.go @@ -80,3 +80,43 @@ func WithTimeout(parent Context, timeout time.Duration) (Context, context.Cancel parent.Context = new return parent, cancel } + +// Map the context to a map +func (ctx *Context) Map() map[string]interface{} { + data := map[string]interface{}{ + "sid": ctx.Sid, + } + + if ctx.ChatID != "" { + data["chat_id"] = ctx.ChatID + } + if ctx.AssistantID != "" { + data["assistant_id"] = ctx.AssistantID + } + if ctx.Stack != "" { + data["stack"] = ctx.Stack + } + if ctx.Path != "" { + data["pathname"] = ctx.Path + } + if len(ctx.FormData) > 0 { + data["formdata"] = ctx.FormData + } + if ctx.Field != nil { + data["field"] = ctx.Field + } + if ctx.Namespace != "" { + data["namespace"] = ctx.Namespace + } + if len(ctx.Config) > 0 { + data["config"] = ctx.Config + } + if ctx.Signal != nil { + data["signal"] = ctx.Signal + } + if ctx.Upload != nil { + data["upload"] = ctx.Upload + } + + return data +} diff --git a/neo/neo.go b/neo/neo.go index f374bc1d..884f4a68 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -25,19 +25,30 @@ func (neo *DSL) Answer(ctx chatctx.Context, question string, c *gin.Context) err return err } - // Get the assistant_id, chat_id - res, err := neo.HookCreate(ctx, messages, c) + var res *assistant.ResHookInit = nil + var ast assistant.API = neo.Assistant + + if ctx.AssistantID != "" { + ast, err = neo.Select(ctx.AssistantID) + if err != nil { + return err + } + } + + // Init the assistant + res, err = ast.HookInit(ctx, []message.Message{{Text: question}}) if err != nil { - msg := message.New().Error(err).Done() - msg.Write(c.Writer) return err } - // Select Assistant - ast, err := neo.Select(res.AssistantID) - if err != nil { - return err + // Switch to the new assistant if necessary + if res.AssistantID != ctx.AssistantID { + ast, err = neo.Select(res.AssistantID) + if err != nil { + return err + } } + // Chat with AI return neo.chat(ast, ctx, messages, c) }