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

@@ -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...)
+}

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

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
+}

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)
 }