From 830bbf7a0efd1fc1ebf66eae82cc948c51115e81 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 19 Jan 2025 15:19:08 +0800 Subject: [PATCH] Refactor Neo API assistant message handling and data structures - Updated message appending methods to use AppendTo for improved clarity and consistency. - Changed output types in HookStream and HookDone methods from string to []message.Data, enhancing data handling capabilities. - Introduced a new Map method in the Data struct for better representation of message data. - Streamlined JSON marshaling in the Data struct to ensure accurate data serialization. These changes enhance the maintainability and robustness of the Neo API assistant, paving the way for improved message handling and data processing. --- neo/assistant/api.go | 19 ++++++--------- neo/assistant/hooks.go | 54 ++++++++++++++++++++++++++++++++++++----- neo/assistant/types.go | 8 +++--- neo/message/contents.go | 30 ++++++++++++++++++++++- neo/message/message.go | 4 +-- neo/neo.go | 2 +- 6 files changed, 92 insertions(+), 25 deletions(-) diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 2b5611d6..8ca92147 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -232,15 +232,12 @@ func (ast *Assistant) streamChat( } // Append content and send message - msg.Append(contents) + msg.AppendTo(contents) value := msg.String() if value != "" { // Handle stream - res, err := ast.HookStream(c, ctx, messages, contents.JSON()) + res, err := ast.HookStream(c, ctx, messages, contents.Data) if err == nil && res != nil { - if res.Output != "" { - value = res.Output - } if res.Next != nil { err = res.Next.Execute(c, ctx) @@ -270,16 +267,16 @@ func (ast *Assistant) streamChat( // Complete the stream if msg.IsDone { - if value == "" { - msg.Write(c.Writer) - } + // if value == "" { + // msg.Write(c.Writer) + // } - res, hookErr := ast.HookDone(c, ctx, messages, contents.JSON()) + res, hookErr := ast.HookDone(c, ctx, messages, contents.Data) if hookErr == nil && res != nil { - if res.Output != "" { + if res.Output != nil { chatMessage.New(). Map(map[string]interface{}{ - "text": res.Output, + "text": res.Input, "done": true, }). Write(c.Writer) diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index f422b60b..ed9c32cf 100644 --- a/neo/assistant/hooks.go +++ b/neo/assistant/hooks.go @@ -69,7 +69,7 @@ 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) (*ResHookStream, error) { +func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output []message.Data) (*ResHookStream, error) { // Create timeout context ctx, cancel := ast.createTimeoutContext(c) @@ -87,8 +87,24 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input switch v := v.(type) { case map[string]interface{}: if res, ok := v["output"].(string); ok { - response.Output = res + vv := []message.Data{} + err := jsoniter.UnmarshalFromString(res, &vv) + if err != nil { + return nil, err + } + response.Output = vv } + + if res, ok := v["output"].([]interface{}); ok { + vv := []message.Data{} + raw, _ := jsoniter.MarshalToString(res) + err := jsoniter.UnmarshalFromString(raw, &vv) + if err != nil { + return nil, err + } + response.Output = vv + } + if res, ok := v["next"].(map[string]interface{}); ok { response.Next = &NextAction{} if name, ok := res["action"].(string); ok { @@ -105,14 +121,19 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input } case string: - response.Output = v + vv := []message.Data{} + err := jsoniter.UnmarshalFromString(v, &vv) + if err != nil { + return nil, err + } + response.Output = vv } return response, nil } // HookDone Handle completion of assistant response -func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookDone, error) { +func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output []message.Data) (*ResHookDone, error) { // Create timeout context ctx, cancel := ast.createTimeoutContext(c) defer cancel() @@ -133,8 +154,24 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input [] switch v := v.(type) { case map[string]interface{}: if res, ok := v["output"].(string); ok { - response.Output = res + vv := []message.Data{} + err := jsoniter.UnmarshalFromString(res, &vv) + if err != nil { + return nil, err + } + response.Output = vv } + + if res, ok := v["output"].([]interface{}); ok { + vv := []message.Data{} + raw, _ := jsoniter.MarshalToString(res) + err := jsoniter.UnmarshalFromString(raw, &vv) + if err != nil { + return nil, err + } + response.Output = vv + } + if res, ok := v["next"].(map[string]interface{}); ok { response.Next = &NextAction{} if name, ok := res["action"].(string); ok { @@ -145,7 +182,12 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input [] } } case string: - response.Output = v + vv := []message.Data{} + err := jsoniter.UnmarshalFromString(v, &vv) + if err != nil { + return nil, err + } + response.Output = vv } return response, nil diff --git a/neo/assistant/types.go b/neo/assistant/types.go index eb6ee90f..f18ad028 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -39,16 +39,16 @@ type ResHookInit struct { // ResHookStream the response of the stream hook type ResHookStream struct { - Silent bool `json:"silent,omitempty"` // Whether to suppress the output - Next *NextAction `json:"next,omitempty"` // The next action - Output string `json:"output,omitempty"` // The output + Silent bool `json:"silent,omitempty"` // Whether to suppress the output + Next *NextAction `json:"next,omitempty"` // The next action + Output []message.Data `json:"output,omitempty"` // The output } // ResHookDone the response of the done hook type ResHookDone struct { Next *NextAction `json:"next,omitempty"` Input []message.Message `json:"input,omitempty"` - Output string `json:"output,omitempty"` + Output []message.Data `json:"output,omitempty"` } // ResHookFail the response of the fail hook diff --git a/neo/message/contents.go b/neo/message/contents.go index 62ed6851..f7a6440a 100644 --- a/neo/message/contents.go +++ b/neo/message/contents.go @@ -120,6 +120,34 @@ func (c *Contents) Text() string { return string(c.Data[c.Current].Bytes) } +// Map returns the map representation +func (data *Data) Map() (map[string]interface{}, error) { + v := map[string]interface{}{"type": data.Type} + + if data.ID != "" { + v["id"] = data.ID + } + + if data.Bytes != nil { + v["text"] = 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 v, nil +} + // MarshalJSON returns the json representation func (data *Data) MarshalJSON() ([]byte, error) { @@ -130,7 +158,7 @@ func (data *Data) MarshalJSON() ([]byte, error) { } if data.Bytes != nil { - v["bytes"] = string(data.Bytes) + v["text"] = string(data.Bytes) } if data.Arguments != nil { diff --git a/neo/message/message.go b/neo/message/message.go index d96f82e4..815d5ea3 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -184,8 +184,8 @@ func (m *Message) SetContent(content string) *Message { return m } -// Append append the contents -func (m *Message) Append(contents *Contents) *Message { +// AppendTo append the contents +func (m *Message) AppendTo(contents *Contents) *Message { switch m.Type { case "text": diff --git a/neo/neo.go b/neo/neo.go index e34a55b7..89cae336 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -126,7 +126,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st } // Append content and send message - msg.Append(contents) + msg.AppendTo(contents) if !silent { value := msg.String() if value != "" {