diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 21ab6b14..b42138dd 100644 --- a/neo/assistant/api.go +++ b/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 } diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index 97de8222..5fa42713 100644 --- a/neo/assistant/hooks.go +++ b/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) - defer cancel() - + ctx := ast.createBackgroundContext() 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{ - Input: input, - Output: contents.Data, - } + response := &ResHookDone{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) { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) +func (ast *Assistant) createTimeoutContext(time time.Duration) (context.Context, context.CancelFunc) { + ctx, cancel := context.WithTimeout(context.Background(), time) return ctx, cancel } diff --git a/neo/hooks.go b/neo/hooks.go index 9a3b6e58..a777f768 100644 --- a/neo/hooks.go +++ b/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 diff --git a/neo/types.go b/neo/types.go index bbf614d2..e02de507 100644 --- a/neo/types.go +++ b/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"` - ChatID string `json:"chat_id,omitempty"` + AssistantID string `json:"assistant_id,omitempty"` + ChatID string `json:"chat_id,omitempty"` + Input []message.Message `json:"messages,omitempty"` }