Merge pull request #839 from trheyi/main

Enhance Neo API assistant with improved message handling and hook sup…

neo/assistant/api.go

@@ -46,12 +46,21 @@ func GetByConnector(connector string, name string) (*Assistant, error) {
 
 // Execute implements the execute functionality
 func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error {
+	contents := chatMessage.NewContents()
+	return ast.execute(c, ctx, input, options, contents)
+}
+
+// Execute implements the execute functionality
+func (ast *Assistant) execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}, contents *chatMessage.Contents) error {
+
 	messages, err := ast.withHistory(ctx, input)
 	if err != nil {
 		return err
 	}
 
-	contents := chatMessage.NewContents()
+	if contents == nil {
+		contents = chatMessage.NewContents()
+	}
 	options = ast.withOptions(options)
 
 	// Add RAG and Version support
@@ -85,7 +94,7 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
 
 	// Handle next action
 	if res != nil && res.Next != nil {
-		return res.Next.Execute(c, ctx)
+		return res.Next.Execute(c, ctx, contents)
 	}
 
 	// Update options if provided
@@ -103,7 +112,7 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
 }
 
 // Execute the next action
-func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error {
+func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context, contents *chatMessage.Contents) error {
 	switch next.Action {
 
 	case "process":
@@ -164,7 +173,7 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error {
 		if v, ok := next.Payload["options"].(map[string]interface{}); ok {
 			options = v
 		}
-		return assistant.Execute(c, ctx, input, options)
+		return assistant.execute(c, ctx, input, options, contents)
 
 	case "exit":
 		return nil
@@ -264,11 +273,11 @@ func (ast *Assistant) streamChat(
 			value := msg.String()
 			if value != "" {
 				// Handle stream
-				res, err := ast.HookStream(c, ctx, messages, contents)
+				res, err := ast.HookStream(c, ctx, messages, msg, contents)
 				if err == nil && res != nil {
 
 					if res.Next != nil {
-						err = res.Next.Execute(c, ctx)
+						err = res.Next.Execute(c, ctx, contents)
 						if err != nil {
 							chatMessage.New().Error(err.Error()).Done().Write(c.Writer)
 						}
@@ -299,22 +308,17 @@ func (ast *Assistant) streamChat(
 				// 	msg.Write(c.Writer)
 				// }
 
+				fmt.Println("Done", contents.JSON())
+
 				res, hookErr := ast.HookDone(c, ctx, messages, contents)
 				if hookErr == nil && res != nil {
-					if res.Output != nil {
-						chatMessage.New().
-							Map(map[string]interface{}{
-								"text": res.Input,
-								"done": true,
-							}).
-							Write(c.Writer)
-					}
 
 					if res.Next != nil {
-						err := res.Next.Execute(c, ctx)
+						err := res.Next.Execute(c, ctx, contents)
 						if err != nil {
 							chatMessage.New().Error(err.Error()).Done().Write(c.Writer)
 						}
+
 						done <- true
 						return 0 // break
 					}
@@ -331,6 +335,16 @@ func (ast *Assistant) streamChat(
 						Write(c.Writer)
 				}
 
+				// Output
+				if res.Output != nil {
+					chatMessage.New().
+						Map(map[string]interface{}{
+							"text": res.Input,
+							"done": true,
+						}).
+						Write(c.Writer)
+				}
+
 				done <- true
 				return 0 // break
 			}

neo/assistant/hooks.go

@@ -17,9 +17,7 @@ import (
 // HookInit initialize the assistant
 func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []message.Message, options map[string]interface{}, contents *message.Contents) (*ResHookInit, error) {
 	// Create timeout context
-	ctx, cancel := ast.createTimeoutContext(c)
+	ctx := ast.createBackgroundContext()
-	defer cancel()
-
 	v, err := ast.call(ctx, "Init", c, contents, context, input, options)
 	if err != nil {
 		if err.Error() == HookErrorMethodNotFound {
@@ -72,13 +70,13 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
 }
 
 // HookStream Handle streaming response from LLM
-func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, contents *chatMessage.Contents) (*ResHookStream, error) {
+func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, msg *message.Message, contents *chatMessage.Contents) (*ResHookStream, error) {
 
 	// Create timeout context
-	ctx, cancel := ast.createTimeoutContext(c)
+	ctx, cancel := ast.createTimeoutContext(5 * time.Second)
 	defer cancel()
 
-	v, err := ast.call(ctx, "Stream", c, contents, context, input)
+	v, err := ast.call(ctx, "Stream", c, contents, context, input, msg, contents.JSON())
 	if err != nil {
 		if err.Error() == HookErrorMethodNotFound {
 			return nil, nil
@@ -140,7 +138,7 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
 	// Create timeout context
 	ctx := ast.createBackgroundContext()
 
-	v, err := ast.call(ctx, "Done", c, contents, context, input)
+	v, err := ast.call(ctx, "Done", c, contents, context, input, contents.Data)
 	if err != nil {
 		if err.Error() == HookErrorMethodNotFound {
 			return nil, nil
@@ -148,10 +146,7 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
 		return nil, err
 	}
 
-	response := &ResHookDone{
+	response := &ResHookDone{Input: input, Output: contents.Data}
-		Input:  input,
-		Output: contents.Data,
-	}
 
 	switch v := v.(type) {
 	case map[string]interface{}:
@@ -198,7 +193,7 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
 // HookFail Handle failure of assistant response
 func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []message.Message, err error, contents *chatMessage.Contents) (*ResHookFail, error) {
 	// Create timeout context
-	ctx, cancel := ast.createTimeoutContext(c)
+	ctx, cancel := ast.createTimeoutContext(5 * time.Second)
 	defer cancel()
 
 	v, callErr := ast.call(ctx, "Fail", c, contents, context, input, err.Error())
@@ -240,8 +235,8 @@ func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []
 }
 
 // createTimeoutContext creates a timeout context with 5 seconds timeout
-func (ast *Assistant) createTimeoutContext(c *gin.Context) (context.Context, context.CancelFunc) {
+func (ast *Assistant) createTimeoutContext(time time.Duration) (context.Context, context.CancelFunc) {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), time)
 	return ctx, cancel
 }
 

neo/hooks.go

@@ -8,6 +8,7 @@ import (
 	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/process"
 	chatctx "github.com/yaoapp/yao/neo/context"
+	"github.com/yaoapp/yao/neo/message"
 )
 
 // HookCreate create the assistant
@@ -25,7 +26,7 @@ func (neo *DSL) HookCreate(ctx chatctx.Context, messages []map[string]interface{
 	}
 
 	// Create a context with 10 second timeout
-	timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
 	defer cancel()
 
 	p, err := process.Of(neo.Create, ctx, messages, c.Writer)
@@ -63,10 +64,23 @@ func (neo *DSL) HookCreate(ctx chatctx.Context, messages []map[string]interface{
 			chatID = ctx.ChatID
 		}
 
-		return CreateResponse{AssistantID: assistantID, ChatID: chatID}, nil
+		// Messages fixed input
+		input := []message.Message{}
+		if vv, has := v["input"]; has {
+			bytes, err := jsoniter.Marshal(vv)
+			if err != nil {
+				return CreateResponse{}, err
+			}
+			err = jsoniter.Unmarshal(bytes, &input)
+			if err != nil {
+				return CreateResponse{}, err
+			}
+		}
+
+		return CreateResponse{AssistantID: assistantID, ChatID: chatID, Input: input}, nil
 	}
 
-	return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID}, nil
+	return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID, Input: nil}, nil
 }
 
 // HookPrepare executes the prepare hook before AI is called

neo/types.go

@@ -3,6 +3,7 @@ package neo
 import (
 	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/yao/neo/assistant"
+	"github.com/yaoapp/yao/neo/message"
 	"github.com/yaoapp/yao/neo/rag"
 	"github.com/yaoapp/yao/neo/store"
 	"github.com/yaoapp/yao/neo/vision"
@@ -62,6 +63,7 @@ type FileUpload struct {
 
 // CreateResponse the response of the create hook
 type CreateResponse struct {
-	AssistantID string `json:"assistant_id,omitempty"`
+	AssistantID string            `json:"assistant_id,omitempty"`
-	ChatID      string `json:"chat_id,omitempty"`
+	ChatID      string            `json:"chat_id,omitempty"`
+	Input       []message.Message `json:"messages,omitempty"`
 }