diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 41f2b0d4..2b5611d6 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -56,6 +56,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, // Run init hook res, err := ast.HookInit(c, ctx, messages, options) if err != nil { + chatMessage.New(). + Assistant(ast.ID, ast.Name, ast.Avatar). + Error(err). + Done(). + Write(c.Writer) return err } @@ -63,6 +68,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, if res != nil && res.AssistantID != ctx.AssistantID { newAst, err := Get(res.AssistantID) if err != nil { + chatMessage.New(). + Assistant(ast.ID, ast.Name, ast.Avatar). + Error(err). + Done(). + Write(c.Writer) return err } *ast = *newAst @@ -163,16 +173,16 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error { func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []chatMessage.Message, options map[string]interface{}) error { clientBreak := make(chan bool, 1) done := make(chan bool, 1) - content := chatMessage.NewContent("text") + contents := chatMessage.NewContents() // Chat with AI in background go func() { - err := ast.streamChat(c, ctx, messages, options, clientBreak, done, content) + err := ast.streamChat(c, ctx, messages, options, clientBreak, done, contents) if err != nil { chatMessage.New().Error(err).Done().Write(c.Writer) } - ast.saveChatHistory(ctx, messages, content) + ast.saveChatHistory(ctx, messages, contents) done <- true }() @@ -194,7 +204,7 @@ func (ast *Assistant) streamChat( options map[string]interface{}, clientBreak chan bool, done chan bool, - content *chatMessage.Content) error { + contents *chatMessage.Contents) error { return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int { select { @@ -210,7 +220,7 @@ func (ast *Assistant) streamChat( // Handle error if msg.Type == "error" { value := msg.String() - res, hookErr := ast.HookFail(c, ctx, messages, content.String(), fmt.Errorf("%s", value)) + res, hookErr := ast.HookFail(c, ctx, messages, contents.JSON(), fmt.Errorf("%s", value)) if hookErr == nil && res != nil && (res.Output != "" || res.Error != "") { value = res.Output if res.Error != "" { @@ -221,26 +231,12 @@ func (ast *Assistant) streamChat( return 0 // break } - // Handle tool call - if msg.Type == "tool_calls" { - content.SetType("function") // Set type to function - // Set id - if id, ok := msg.Props["id"].(string); ok && id != "" { - content.SetID(id) - } - - // Set name - if name, ok := msg.Props["name"].(string); ok && name != "" { - content.SetName(name) - } - } - // Append content and send message + msg.Append(contents) value := msg.String() - content.Append(value) if value != "" { // Handle stream - res, err := ast.HookStream(c, ctx, messages, content.String(), content.Type == "function") + res, err := ast.HookStream(c, ctx, messages, contents.JSON()) if err == nil && res != nil { if res.Output != "" { value = res.Output @@ -278,9 +274,7 @@ func (ast *Assistant) streamChat( msg.Write(c.Writer) } - // Call HookDone - content.SetStatus(chatMessage.ContentStatusDone) - res, hookErr := ast.HookDone(c, ctx, messages, content.String(), content.Type == "function") + res, hookErr := ast.HookDone(c, ctx, messages, contents.JSON()) if hookErr == nil && res != nil { if res.Output != "" { chatMessage.New(). @@ -322,8 +316,8 @@ func (ast *Assistant) streamChat( } // saveChatHistory saves the chat history if storage is available -func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, content *chatMessage.Content) { - if len(content.Bytes) > 0 && ctx.Sid != "" && len(messages) > 0 { +func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, contents *chatMessage.Contents) { + if len(contents.Data) > 0 && ctx.Sid != "" && len(messages) > 0 { userMessage := messages[len(messages)-1] data := []map[string]interface{}{ { @@ -333,7 +327,7 @@ func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessag }, { "role": "assistant", - "content": content.String(), + "content": contents.JSON(), "name": ctx.Sid, "assistant_id": ast.ID, "assistant_name": ast.Name, diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index bf21fa25..f422b60b 100644 --- a/neo/assistant/hooks.go +++ b/neo/assistant/hooks.go @@ -6,6 +6,7 @@ import ( "time" "github.com/gin-gonic/gin" + jsoniter "github.com/json-iterator/go" chatctx "github.com/yaoapp/yao/neo/context" "github.com/yaoapp/yao/neo/message" ) @@ -34,6 +35,17 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input [] response.ChatID = res } + // input + if input, has := v["input"]; has { + raw, _ := jsoniter.MarshalToString(input) + vv := []message.Message{} + err := jsoniter.UnmarshalFromString(raw, &vv) + if err != nil { + return nil, err + } + response.Input = vv + } + if res, ok := v["next"].(map[string]interface{}); ok { response.Next = &NextAction{} if name, ok := res["action"].(string); ok { @@ -57,13 +69,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, output string, toolcall bool) (*ResHookStream, error) { +func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookStream, error) { // Create timeout context ctx, cancel := ast.createTimeoutContext(c) defer cancel() - v, err := ast.call(ctx, "Stream", context, input, output, toolcall, c.Writer) + v, err := ast.call(ctx, "Stream", context, input, output, c.Writer) if err != nil { if err.Error() == HookErrorMethodNotFound { return nil, nil @@ -100,12 +112,12 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input } // HookDone Handle completion of assistant response -func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string, toolcall bool) (*ResHookDone, error) { +func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookDone, error) { // Create timeout context ctx, cancel := ast.createTimeoutContext(c) defer cancel() - v, err := ast.call(ctx, "Done", context, input, output, toolcall, c.Writer) + v, err := ast.call(ctx, "Done", context, input, output, c.Writer) if err != nil { if err.Error() == HookErrorMethodNotFound { return nil, nil @@ -185,7 +197,7 @@ 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(c.Request.Context(), 5*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) return ctx, cancel } diff --git a/neo/assistant/load.go b/neo/assistant/load.go index 9cbd7dff..13563acb 100644 --- a/neo/assistant/load.go +++ b/neo/assistant/load.go @@ -28,6 +28,10 @@ var defaultConnector string = "" // default connector // LoadBuiltIn load the built-in assistants func LoadBuiltIn() error { + + // Clear the cache + loaded.Clear() + root := `/assistants` app, err := fs.Get("app") if err != nil { diff --git a/neo/message/content.go b/neo/message/content.go deleted file mode 100644 index 2d2466c5..00000000 --- a/neo/message/content.go +++ /dev/null @@ -1,100 +0,0 @@ -package message - -import ( - jsoniter "github.com/json-iterator/go" -) - -const ( - // ContentStatusPending the content status pending - ContentStatusPending = iota - // ContentStatusDone the content status done - ContentStatusDone - // ContentStatusError the content status error - ContentStatusError -) - -// Content the content -type Content struct { - ID string `json:"id"` - Name string `json:"name"` - Bytes []byte `json:"bytes"` - Type string `json:"type"` // text, function, error - Status uint8 `json:"status"` // 0: pending, 1: done -} - -// NewContent create a new content -func NewContent(typ string) *Content { - if typ == "" { - typ = "text" - } - - return &Content{ - Bytes: []byte{}, - Type: typ, - Status: ContentStatusPending, - } -} - -// String the content string -func (c *Content) String() string { - if c.Type == "function" { - - var arguments interface{} = string(c.Bytes) - if c.Status == ContentStatusDone { - var vv interface{} = nil - err := jsoniter.Unmarshal(c.Bytes, &vv) - if err != nil { - return "" - } - arguments = vv - } - - data := map[string]interface{}{ - "id": c.ID, - "type": "function", - "text": c.Name, - "props": map[string]interface{}{ - "id": c.ID, - "name": c.Name, - "arguments": arguments, - }, - "function": map[string]interface{}{ - "name": c.Name, - "arguments": arguments, - }, - } - - raw, err := jsoniter.MarshalToString(data) - if err != nil { - return "" - } - - return raw - } - return string(c.Bytes) -} - -// SetID set the content id -func (c *Content) SetID(id string) { - c.ID = id -} - -// SetName set the content name -func (c *Content) SetName(name string) { - c.Name = name -} - -// SetType set the content type -func (c *Content) SetType(typ string) { - c.Type = typ -} - -// Append append the content -func (c *Content) Append(data string) { - c.Bytes = append(c.Bytes, []byte(data)...) -} - -// SetStatus set the content status -func (c *Content) SetStatus(status uint8) { - c.Status = status -} diff --git a/neo/message/contents.go b/neo/message/contents.go new file mode 100644 index 00000000..62ed6851 --- /dev/null +++ b/neo/message/contents.go @@ -0,0 +1,150 @@ +package message + +import ( + jsoniter "github.com/json-iterator/go" +) + +const ( + // ContentStatusPending the content status pending + ContentStatusPending = iota + // ContentStatusDone the content status done + ContentStatusDone + // ContentStatusError the content status error + ContentStatusError +) + +// Contents the contents +type Contents struct { + Current int `json:"current"` // the current content index + Data []Data `json:"data"` // the data +} + +// Data the data of the content +type Data struct { + Type string `json:"type"` // text, function, error, ... + ID string `json:"id"` // the id of the content + Function string `json:"function"` // the function name + Bytes []byte `json:"bytes"` // the content bytes + Arguments []byte `json:"arguments"` // the function arguments +} + +// NewContents create a new contents +func NewContents() *Contents { + return &Contents{ + Current: -1, + Data: []Data{}, + } +} + +// NewText create a new text data and append to the contents +func (c *Contents) NewText(bytes []byte) *Contents { + c.Data = append(c.Data, Data{ + Type: "text", + Bytes: bytes, + }) + c.Current++ + return c +} + +// NewFunction create a new function data and append to the contents +func (c *Contents) NewFunction(function string, arguments []byte) *Contents { + c.Data = append(c.Data, Data{ + Type: "function", + Function: function, + Arguments: arguments, + }) + c.Current++ + return c +} + +// SetFunctionID set the id of the current function content +func (c *Contents) SetFunctionID(id string) *Contents { + if c.Current == -1 { + c.NewFunction("", []byte{}) + } + c.Data[c.Current].ID = id + return c +} + +// NewError create a new error data and append to the contents +func (c *Contents) NewError(err []byte) *Contents { + c.Data = append(c.Data, Data{ + Type: "error", + Bytes: err, + }) + c.Current++ + return c +} + +// AppendText append the text to the current content +func (c *Contents) AppendText(bytes []byte) *Contents { + if c.Current == -1 { + c.NewText(bytes) + return c + } + c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, bytes...) + return c +} + +// AppendFunction append the function to the current content +func (c *Contents) AppendFunction(arguments []byte) *Contents { + if c.Current == -1 { + c.NewFunction("", arguments) + return c + } + c.Data[c.Current].Arguments = append(c.Data[c.Current].Arguments, arguments...) + return c +} + +// AppendError append the error to the current content +func (c *Contents) AppendError(err []byte) *Contents { + if c.Current == -1 { + c.NewError(err) + return c + } + c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, err...) + return c +} + +// JSON returns the json representation +func (c *Contents) JSON() string { + raw, _ := jsoniter.MarshalToString(c.Data) + return raw +} + +// Text returns the text of the current content +func (c *Contents) Text() string { + if c.Current == -1 { + return "" + } + return string(c.Data[c.Current].Bytes) +} + +// MarshalJSON returns the json representation +func (data *Data) MarshalJSON() ([]byte, error) { + + v := map[string]interface{}{"type": data.Type} + + if data.ID != "" { + v["id"] = data.ID + } + + if data.Bytes != nil { + v["bytes"] = string(data.Bytes) + } + + if data.Arguments != nil { + var vv interface{} = nil + err := jsoniter.Unmarshal(data.Arguments, &vv) + if err != nil { + return nil, err + } + v["arguments"] = vv + } + + if data.Function != "" { + v["function"] = data.Function + } + + return jsoniter.Marshal(v) +} diff --git a/neo/message/message.go b/neo/message/message.go index 39551102..d96f82e4 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -99,7 +99,7 @@ func NewOpenAI(data []byte) *Message { msg.Type = "tool_calls" if len(toolCalls.Choices) > 0 && len(toolCalls.Choices[0].Delta.ToolCalls) > 0 { msg.Props["id"] = toolCalls.Choices[0].Delta.ToolCalls[0].ID - msg.Props["name"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name + msg.Props["function"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name msg.Text = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments } @@ -184,6 +184,32 @@ func (m *Message) SetContent(content string) *Message { return m } +// Append append the contents +func (m *Message) Append(contents *Contents) *Message { + + switch m.Type { + case "text": + if m.Text != "" { + contents.AppendText([]byte(m.Text)) + } + + case "tool_calls": + + // Set function name + if name, ok := m.Props["function"].(string); ok && name != "" { + contents.NewFunction(name, []byte(m.Text)) + } + + // Set id + if id, ok := m.Props["id"].(string); ok && id != "" { + contents.SetFunctionID(id) + } + + contents.AppendFunction([]byte(m.Text)) + } + return m +} + // Content get the content func (m *Message) Content() string { content := map[string]interface{}{"text": m.Text} @@ -291,6 +317,14 @@ func (m *Message) Done() *Message { return m } +// Assistant set the assistant +func (m *Message) Assistant(id string, name string, avatar string) *Message { + m.AssistantID = id + m.AssistantName = name + m.AssistantAvatar = avatar + return m +} + // Action add an action func (m *Message) Action(name string, t string, payload interface{}, next string) *Message { if m.Data != nil { @@ -339,11 +373,6 @@ func (m *Message) Write(w gin.ResponseWriter) bool { return true } -// Append appends content to the byte slice -func (m *Message) Append(content []byte) []byte { - return append(content, []byte(m.Text)...) -} - // WriteError writes an error message to response writer func (m *Message) WriteError(w gin.ResponseWriter, message string) { errMsg := strings.Trim(exception.New(message, 500).Message, "\"") diff --git a/neo/neo.go b/neo/neo.go index c9a45f98..e34a55b7 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -100,7 +100,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st clientBreak := make(chan bool, 1) done := make(chan bool, 1) fail := make(chan error, 1) - content := []byte{} + contents := message.NewContents() // Chat with AI in background go func() { @@ -126,7 +126,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st } // Append content and send message - content = msg.Append(content) + msg.Append(contents) if !silent { value := msg.String() if value != "" { @@ -166,7 +166,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st // Wait for completion or client disconnect select { case <-done: - return string(content), nil + return contents.Text(), nil case err := <-fail: return "", err case <-c.Writer.CloseNotify():