diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 685e1bf2..bc084096 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -6,7 +6,9 @@ import ( "fmt" "strings" + "github.com/gin-gonic/gin" "github.com/yaoapp/gou/fs" + chatctx "github.com/yaoapp/yao/neo/context" "github.com/yaoapp/yao/neo/message" chatMessage "github.com/yaoapp/yao/neo/message" ) @@ -41,6 +43,184 @@ func GetByConnector(connector string, name string) (*Assistant, error) { return assistant, nil } +// Execute implements the execute functionality +func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error { + messages, err := ast.withHistory(ctx, input) + if err != nil { + return err + } + + options = ast.withOptions(options) + + // Run init hook + res, err := ast.HookInit(c, ctx, messages, options) + if err != nil { + return err + } + + // Switch to the new assistant if necessary + if res.AssistantID != ctx.AssistantID { + newAst, err := Get(res.AssistantID) + if err != nil { + return err + } + *ast = *newAst + } + + // Handle next action + if res.Next != nil { + switch res.Next.Action { + case "exit": + return nil + // Add other actions here if needed + } + } + + // Update options if provided + if res.Options != nil { + options = res.Options + } + + // Only proceed with chat stream if no specific next action was handled + return ast.handleChatStream(c, ctx, messages, options) +} + +// handleChatStream manages the streaming chat interaction with the AI +func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []message.Message, options map[string]interface{}) error { + clientBreak := make(chan bool, 1) + done := make(chan bool, 1) + content := []byte{} + + // Chat with AI in background + go func() { + err := ast.streamChat(c, messages, options, clientBreak, done, &content) + if err != nil { + chatMessage.New().Error(err).Done().Write(c.Writer) + } + + ast.saveChatHistory(ctx, messages, content) + done <- true + }() + + // Wait for completion or client disconnect + select { + case <-done: + return nil + case <-c.Writer.CloseNotify(): + clientBreak <- true + return nil + } +} + +// streamChat handles the streaming chat interaction +func (ast *Assistant) streamChat(c *gin.Context, messages []message.Message, options map[string]interface{}, + clientBreak chan bool, done chan bool, content *[]byte) error { + + return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int { + select { + case <-clientBreak: + return 0 // break + + default: + msg := chatMessage.NewOpenAI(data) + if msg == nil { + return 1 // continue + } + + // Handle error + if msg.Type == "error" { + value := msg.String() + chatMessage.New().Error(value).Done().Write(c.Writer) + return 0 // break + } + + // Append content and send message + *content = msg.Append(*content) + value := msg.String() + if value != "" { + chatMessage.New(). + Map(map[string]interface{}{ + "text": value, + "done": msg.IsDone, + }). + Write(c.Writer) + } + + // Complete the stream + if msg.IsDone { + if value == "" { + msg.Write(c.Writer) + } + done <- true + return 0 // break + } + + return 1 // continue + } + }) +} + +// saveChatHistory saves the chat history if storage is available +func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []message.Message, content []byte) { + if len(content) > 0 && ctx.Sid != "" && len(messages) > 0 { + storage.SaveHistory( + ctx.Sid, + []map[string]interface{}{ + {"role": "user", "content": messages[len(messages)-1].Content(), "name": ctx.Sid}, + {"role": "assistant", "content": string(content), "name": ctx.Sid}, + }, + ctx.ChatID, + nil, + ) + } +} + +func (ast *Assistant) withOptions(options map[string]interface{}) map[string]interface{} { + if options == nil { + options = map[string]interface{}{} + } + + if ast.Options != nil { + for key, value := range ast.Options { + options[key] = value + } + } + return options +} + +func (ast *Assistant) withPrompts(messages []message.Message) []message.Message { + if ast.Prompts != nil { + for _, prompt := range ast.Prompts { + name := ast.Name + if prompt.Name != "" { + name = prompt.Name + } + messages = append(messages, *message.New().Map(map[string]interface{}{"role": prompt.Role, "content": prompt.Content, "name": name})) + } + } + return messages +} + +func (ast *Assistant) withHistory(ctx chatctx.Context, input string) ([]message.Message, error) { + messages := []message.Message{} + messages = ast.withPrompts(messages) + if storage != nil { + history, err := storage.GetHistory(ctx.Sid, ctx.ChatID) + if err != nil { + return nil, err + } + + // Add history messages + for _, h := range history { + messages = append(messages, *message.New().Map(h)) + } + } + + // Add user message + messages = append(messages, *message.New().Map(map[string]interface{}{"role": "user", "content": input, "name": ctx.Sid})) + return messages, nil +} + // Chat implements the chat functionality func (ast *Assistant) Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error { if ast.openai == nil { diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index 57f1329b..bbe8894f 100644 --- a/neo/assistant/hooks.go +++ b/neo/assistant/hooks.go @@ -17,17 +17,26 @@ const ( // ResHookInit the response of the init hook type ResHookInit 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"` + Next *NextAction `json:"next,omitempty"` + Input []message.Message `json:"input,omitempty"` + Options map[string]interface{} `json:"options,omitempty"` +} + +// NextAction the next action +type NextAction struct { + Action string `json:"action"` + Payload map[string]interface{} `json:"payload,omitempty"` } // HookInit initialize the assistant -func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, messages []message.Message) (*ResHookInit, error) { +func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error) { // Create timeout context ctx, cancel := ast.createTimeoutContext(c) defer cancel() - v, err := ast.call(ctx, "Init", context, messages, c.Writer) + v, err := ast.call(ctx, "Init", context, input, c.Writer) if err != nil { if err.Error() == HookErrorMethodNotFound { return nil, nil diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 54bc47f8..9877b11e 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -19,7 +19,8 @@ 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(c *gin.Context, ctx chatctx.Context, messages []message.Message) (*ResHookInit, error) + Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error + HookInit(c *gin.Context, ctx chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error) } // RAG the RAG interface diff --git a/neo/neo.go b/neo/neo.go index 83bd91b2..c9a45f98 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -4,7 +4,6 @@ import ( "fmt" "os" "strings" - "sync" "github.com/gin-gonic/gin" "github.com/yaoapp/kun/log" @@ -13,44 +12,17 @@ import ( "github.com/yaoapp/yao/neo/message" ) -// Lock the assistant list -var lock sync.Mutex = sync.Mutex{} - // Answer reply the message func (neo *DSL) Answer(ctx chatctx.Context, question string, c *gin.Context) error { - messages, err := neo.withHistory(ctx, question) - if err != nil { - msg := message.New().Error(err).Done() - msg.Write(c.Writer) - return err - } - - var res *assistant.ResHookInit = nil + var err error 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(c, ctx, messages) - 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) + return ast.Execute(c, ctx, question, nil) } // Select select an assistant @@ -87,6 +59,7 @@ func (neo *DSL) GenerateChatTitle(ctx chatctx.Context, input string, c *gin.Cont 2. The title should be a single sentence. 3. The title should be in same language as the chat. 4. The title should be no more than 50 characters. + 5. ANSWER ONLY THE TITLE CONTENT, FOR EXAMPLE: Chat with AI is a valid title, but "Chat with AI" is not a valid title. ` isSilent := false if len(silent) > 0 { @@ -273,130 +246,3 @@ func (neo *DSL) Download(ctx chatctx.Context, c *gin.Context) (*assistant.FileRe // Download file using the assistant return ast.Download(ctx.Context, fileID) } - -// chat chat with AI -func (neo *DSL) chat(ast assistant.API, ctx chatctx.Context, messages []message.Message, c *gin.Context) error { - if ast == nil { - msg := message.New().Error("assistant is not initialized").Done() - msg.Write(c.Writer) - return fmt.Errorf("assistant is not initialized") - } - - clientBreak := make(chan bool, 1) - done := make(chan bool, 1) - content := []byte{} - - // Chat with AI in background - go func() { - err := ast.Chat(c.Request.Context(), messages, neo.Option, func(data []byte) int { - select { - case <-clientBreak: - return 0 // break - - default: - msg := message.NewOpenAI(data) - if msg == nil { - return 1 // continue - } - - // Handle error - if msg.Type == "error" { - value := msg.String() - message.New().Error(value).Done().Write(c.Writer) - return 0 // break - } - - // Append content and send message - content = msg.Append(content) - value := msg.String() - if value != "" { - message.New(). - Map(map[string]interface{}{ - "text": value, - "done": msg.IsDone, - }). - Write(c.Writer) - } - - // Complete the stream - if msg.IsDone { - if value == "" { - msg.Write(c.Writer) - } - done <- true - return 0 // break - } - - return 1 // continue - } - }) - - if err != nil { - log.Error("Chat error: %s", err.Error()) - message.New().Error(err).Done().Write(c.Writer) - } - - // Save chat history - if len(content) > 0 { - neo.saveHistory(ctx.Sid, ctx.ChatID, content, messages) - } - - done <- true - }() - - // Wait for completion or client disconnect - select { - case <-done: - return nil - case <-c.Writer.CloseNotify(): - clientBreak <- true - return nil - } -} - -func (neo *DSL) withHistory(ctx chatctx.Context, question string) ([]message.Message, error) { - history, err := neo.Store.GetHistory(ctx.Sid, ctx.ChatID) - if err != nil { - return nil, err - } - - // Add history messages - messages := []message.Message{} - for _, h := range history { - messages = append(messages, *message.New().Map(h)) - } - - // Add user message - messages = append(messages, *message.New().Map(map[string]interface{}{"role": "user", "content": question, "name": ctx.Sid})) - return messages, nil -} - -// saveHistory save the history -func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages []message.Message) { - if len(content) > 0 && sid != "" && len(messages) > 0 { - err := neo.Store.SaveHistory( - sid, - []map[string]interface{}{ - {"role": "user", "content": messages[len(messages)-1].Content(), "name": sid}, - {"role": "assistant", "content": string(content), "name": sid}, - }, - chatID, - nil, - ) - - if err != nil { - log.Error("Save history error: %s", err.Error()) - } - } -} - -// sendMessage sends a message to the client -func (neo *DSL) sendMessage(w gin.ResponseWriter, data interface{}) error { - if msg, ok := data.(map[string]interface{}); ok { - if !message.New().Map(msg).Write(w) { - return fmt.Errorf("failed to write message to stream") - } - return nil - } - return fmt.Errorf("invalid message data type") -}