From 238347c834ea0bb0126c8f61af2fd3d0a82f5145 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 31 Dec 2024 16:01:48 +0800 Subject: [PATCH] Refactor message handling and enhance file reading capabilities in Neo API - Updated the message handling logic to improve error reporting and content management, ensuring more robust communication with clients. - Introduced a new ReadBase64 method in both Local and OpenAI assistant implementations to read files and return their base64 encoded content, enhancing file handling capabilities. - Removed the deprecated JSON message handling code, streamlining the message processing structure. - Refactored message struct to include new fields and methods for better data management and response handling. - Improved the chat functionality to handle attachments and user messages more effectively, ensuring a smoother user experience. --- neo/assistant/local/file.go | 28 ++++ neo/assistant/openai/chat.go | 90 ++++++++++- neo/assistant/openai/file.go | 28 ++++ neo/assistant/types.go | 1 + neo/message/json.go | 266 ------------------------------- neo/message/message.go | 294 ++++++++++++++++++++++++++++++++++- neo/message/types.go | 27 ---- neo/neo.go | 52 ++++--- 8 files changed, 466 insertions(+), 320 deletions(-) delete mode 100644 neo/message/json.go delete mode 100644 neo/message/types.go diff --git a/neo/assistant/local/file.go b/neo/assistant/local/file.go index 79d40b6f..cb7a692a 100644 --- a/neo/assistant/local/file.go +++ b/neo/assistant/local/file.go @@ -3,6 +3,7 @@ package local import ( "context" "crypto/sha256" + "encoding/base64" "fmt" "io" "mime/multipart" @@ -132,3 +133,30 @@ func (ast *Local) Download(ctx context.Context, fileID string) (*assistant.FileR Extension: ext, }, nil } + +// ReadBase64 reads a file and returns its base64 encoded content +func (ast *Local) ReadBase64(ctx context.Context, fileID string) (string, error) { + // Get the data filesystem + data, err := fs.Get("data") + if err != nil { + return "", fmt.Errorf("get filesystem error: %s", err.Error()) + } + + // Check if file exists + exists, err := data.Exists(fileID) + if err != nil { + return "", fmt.Errorf("check file error: %s", err.Error()) + } + if !exists { + return "", fmt.Errorf("file %s not found", fileID) + } + + // Read file content + content, err := data.ReadFile(fileID) + if err != nil { + return "", fmt.Errorf("read file error: %s", err.Error()) + } + + // Encode to base64 + return base64.StdEncoding.EncodeToString(content), nil +} diff --git a/neo/assistant/openai/chat.go b/neo/assistant/openai/chat.go index 765c6f3a..1bf80a8c 100644 --- a/neo/assistant/openai/chat.go +++ b/neo/assistant/openai/chat.go @@ -3,6 +3,9 @@ package openai import ( "context" "fmt" + "strings" + + chatMessage "github.com/yaoapp/yao/neo/message" ) // Chat the chat struct @@ -21,10 +24,95 @@ func (ast *OpenAI) Chat(ctx context.Context, messages []map[string]interface{}, return fmt.Errorf("openai is not initialized") } - _, ext := ast.openai.ChatCompletionsWith(ctx, messages, option, cb) + requestMessages, err := ast.requestMessages(ctx, messages) + if err != nil { + return fmt.Errorf("request messages error: %s", err.Error()) + } + + _, ext := ast.openai.ChatCompletionsWith(ctx, requestMessages, option, cb) if ext != nil { return fmt.Errorf("openai chat completions with error: %s", ext.Message) } return nil } + +func (ast *OpenAI) requestMessages(ctx context.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) { + newMessages := []map[string]interface{}{} + length := len(messages) + for index, message := range messages { + role, ok := message["role"].(string) + if !ok { + return nil, fmt.Errorf("role must be string") + } + + content, ok := message["content"].(string) + if !ok { + return nil, fmt.Errorf("content must be string") + } + + newMessage := map[string]interface{}{ + "role": role, + "content": content, + } + + // Handle name if present + if name, ok := message["name"].(string); ok { + newMessage["name"] = name + } + + newMessage["content"] = content + + // Special handling for user messages with JSON content last message + if role == "user" && index == length-1 { + content = strings.TrimSpace(content) + msg, err := chatMessage.NewString(content) + if err != nil { + return nil, fmt.Errorf("new string error: %s", err.Error()) + } + + newMessage["content"] = msg.Text + if msg.Attachments != nil { + content, err := ast.withAttachments(ctx, msg) + if err != nil { + return nil, fmt.Errorf("with attachments error: %s", err.Error()) + } + newMessage["content"] = content + } + } + + newMessages = append(newMessages, newMessage) + } + return newMessages, nil +} + +func (ast *OpenAI) withAttachments(ctx context.Context, msg *chatMessage.Message) ([]map[string]interface{}, error) { + contents := []map[string]interface{}{{"type": "text", "text": msg.Text}} + images := []string{} + for _, attachment := range msg.Attachments { + if strings.HasPrefix(attachment.ContentType, "image/") { + images = append(images, attachment.FileID) + } + } + + if len(images) == 0 { + return contents, nil + } + + for _, image := range images { + bytes64, err := ast.ReadBase64(ctx, image) + if err != nil { + return nil, fmt.Errorf("read base64 error: %s", err.Error()) + } + + contents = append(contents, map[string]interface{}{ + "type": "image_url", + "image_url": map[string]string{ + "url": fmt.Sprintf("data:image/jpeg;base64,%s", bytes64), + }, + }, + ) + } + + return contents, nil +} diff --git a/neo/assistant/openai/file.go b/neo/assistant/openai/file.go index 44892d49..062a688b 100644 --- a/neo/assistant/openai/file.go +++ b/neo/assistant/openai/file.go @@ -3,6 +3,7 @@ package openai import ( "context" "crypto/sha256" + "encoding/base64" "fmt" "io" "mime/multipart" @@ -137,3 +138,30 @@ func (ast *OpenAI) Download(ctx context.Context, fileID string) (*assistant.File Extension: ext, }, nil } + +// ReadBase64 reads a file and returns its base64 encoded content +func (ast *OpenAI) ReadBase64(ctx context.Context, fileID string) (string, error) { + // Get the data filesystem + data, err := fs.Get("data") + if err != nil { + return "", fmt.Errorf("get filesystem error: %s", err.Error()) + } + + // Check if file exists + exists, err := data.Exists(fileID) + if err != nil { + return "", fmt.Errorf("check file error: %s", err.Error()) + } + if !exists { + return "", fmt.Errorf("file %s not found", fileID) + } + + // Read file content + content, err := data.ReadFile(fileID) + if err != nil { + return "", fmt.Errorf("read file error: %s", err.Error()) + } + + // Encode to base64 + return base64.StdEncoding.EncodeToString(content), nil +} diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 35ad239d..a122ed06 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -11,6 +11,7 @@ type API interface { Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error 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) } // Prompt a prompt diff --git a/neo/message/json.go b/neo/message/json.go deleted file mode 100644 index 1ebda8a1..00000000 --- a/neo/message/json.go +++ /dev/null @@ -1,266 +0,0 @@ -package message - -import ( - "fmt" - "strings" - - "github.com/fatih/color" - "github.com/gin-gonic/gin" - jsoniter "github.com/json-iterator/go" - "github.com/yaoapp/gou/helper" - "github.com/yaoapp/kun/exception" - "github.com/yaoapp/kun/log" - "github.com/yaoapp/kun/maps" - "github.com/yaoapp/yao/openai" -) - -// JSON the JSON message -type JSON struct{ *Message } - -// New create a new JSON message -func New() *JSON { - return &JSON{makeMessage()} -} - -// NewOpenAI create a new JSON message -func NewOpenAI(data []byte) *JSON { - - if data == nil || len(data) == 0 { - return nil - } - - msg := makeMessage() - text := string(data) - data = []byte(strings.TrimPrefix(text, "data: ")) - switch { - case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"content":`): - var message openai.Message - err := jsoniter.Unmarshal(data, &message) - if err != nil { - msg.Text = err.Error() - return &JSON{msg} - } - - if len(message.Choices) > 0 { - msg.Text = message.Choices[0].Delta.Content - } - break - - case strings.Contains(text, `[DONE]`): - msg.Done = true - break - - case strings.Contains(text, `"finish_reason":"stop"`): - msg.Done = true - break - - default: - - str := string(data) - // Remove "data: " and " - str = strings.TrimPrefix(str, "data: ") - str = strings.Trim(str, "\"") - msg.Type = "error" - msg.Text = str - } - - return &JSON{msg} -} - -func (json *JSON) String() string { - if json.Message == nil { - return "" - } - return json.Message.Text -} - -// Text set the text -func (json *JSON) Text(text string) *JSON { - - json.Message.Text = text - if json.Message.Data != nil { - replaced := helper.Bind(text, json.Message.Data) - if replacedText, ok := replaced.(string); ok { - json.Message.Text = replacedText - } - } - - return json -} - -// Error set the error -func (json *JSON) Error(message interface{}) *JSON { - json.Message.Type = "error" - if err, ok := message.(error); ok { - json.Message.Text = err.Error() - } else if msg, ok := message.(string); ok { - json.Message.Text = msg - } else { - json.Message.Text = fmt.Sprintf("%v", message) - } - return json -} - -// Map set from map -func (json *JSON) Map(msg map[string]interface{}) *JSON { - if msg == nil { - return json - } - - if text, ok := msg["text"].(string); ok { - json.Message.Text = text - } - - if typ, ok := msg["type"].(string); ok { - json.Message.Text = typ - } - - if done, ok := msg["done"].(bool); ok { - json.Message.Done = done - } - - if confirm, ok := msg["confirm"].(bool); ok { - json.Message.Confirm = confirm - } - - if command, ok := msg["command"].(map[string]interface{}); ok { - json.Message.Command = &Command{} - if id, ok := command["id"].(string); ok { - json.Message.Command.ID = id - } - if name, ok := command["name"].(string); ok { - json.Message.Command.Name = name - } - if request, ok := command["request"].(string); ok { - json.Message.Command.Reqeust = request - } - } - - if actions, ok := msg["actions"].([]interface{}); ok { - for _, action := range actions { - if v, ok := action.(map[string]interface{}); ok { - action := Action{} - if name, ok := v["name"].(string); ok { - action.Name = name - } - if t, ok := v["type"].(string); ok { - action.Type = t - } - if payload, ok := v["payload"].(map[string]interface{}); ok { - action.Payload = payload - } - - if next, ok := v["next"].(string); ok { - action.Next = next - } - json.Message.Actions = append(json.Message.Actions, action) - } - } - } - - if data, ok := msg["data"].(map[string]interface{}); ok { - json.Message.Data = data - } - - return json -} - -// Done set the done -func (json *JSON) Done() *JSON { - json.Message.Done = true - return json -} - -// Confirm set the confirm -func (json *JSON) Confirm() *JSON { - json.Message.Confirm = true - return json -} - -// Command set the command -func (json *JSON) Command(name, id, request string) *JSON { - json.Message.Command = &Command{ - ID: id, - Name: name, - Reqeust: request, - } - return json -} - -// Action set the action -func (json *JSON) Action(name string, t string, payload interface{}, next string) *JSON { - - if json.Message.Data != nil { - payload = helper.Bind(payload, json.Message.Data) - } - - json.Message.Actions = append(json.Message.Actions, Action{ - Name: name, - Type: t, - Payload: payload, - Next: next, - }) - return json -} - -// Bind replace with data -func (json *JSON) Bind(data map[string]interface{}) *JSON { - if data == nil { - return json - } - - json.Message.Data = maps.Of(data).Dot() - return json -} - -// IsDone check if the message is done -func (json *JSON) IsDone() bool { - return json.Message.Done -} - -// Write the message -func (json *JSON) Write(w gin.ResponseWriter) bool { - - defer func() { - if r := recover(); r != nil { - message := "Write Response Exception: (if clinet close the connection, it's normal) \n %s\n\n" - color.Red(message, r) - } - }() - - data, err := jsoniter.Marshal(json.Message) - if err != nil { - log.Error("%s", err.Error()) - return false - } - - data = append([]byte("data: "), data...) - data = append(data, []byte("\n\n")...) - - _, err = w.Write(data) - if err != nil { - color.Red("Write JSON Message Error: %s", err.Error()) - return false - } - w.Flush() - return true -} - -// Append the message -func (json *JSON) Append(content []byte) []byte { - return append(content, []byte(json.Message.Text)...) -} - -func (json *JSON) writeError(w gin.ResponseWriter, message string) { - data := []byte(`{"text":"` + strings.Trim(exception.New(message, 500).Message, "\"") + `","type":"error"}`) - if json.Message.Done { - data = []byte(`{"text":"` + strings.Trim(exception.New(message, 500).Message, "\"") + `","type":"error", "done":true}`) - } - data = append([]byte("data: "), data...) - data = append(data, []byte("\n\n")...) - _, err := w.Write(data) - if err != nil { - color.Red("Write JSON Message Error: %s", message) - } - w.Flush() -} diff --git a/neo/message/message.go b/neo/message/message.go index 7726dd26..48c3a770 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -1,6 +1,296 @@ package message -// makeMessage create a new message -func makeMessage() *Message { +import ( + "fmt" + "strings" + + "github.com/fatih/color" + "github.com/gin-gonic/gin" + jsoniter "github.com/json-iterator/go" + "github.com/yaoapp/gou/helper" + "github.com/yaoapp/kun/exception" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/kun/maps" + "github.com/yaoapp/yao/openai" +) + +// Message the message +type Message struct { + Text string `json:"text,omitempty"` // text content + Type string `json:"type,omitempty"` // error, text, plan, table, form, page, file, video, audio, image, markdown, json ... + Props map[string]interface{} `json:"props,omitempty"` // props for the types + IsDone bool `json:"done,omitempty"` + Actions []Action `json:"actions,omitempty"` // Conversation Actions for frontend + Attachments []Attachment `json:"attachments,omitempty"` // File attachments + Data map[string]interface{} `json:"-"` +} + +// Attachment represents a file attachment +type Attachment struct { + Name string `json:"name,omitempty"` + URL string `json:"url,omitempty"` + Type string `json:"type,omitempty"` + ContentType string `json:"content_type,omitempty"` + Bytes int64 `json:"bytes,omitempty"` + CreatedAt int64 `json:"created_at,omitempty"` + FileID string `json:"file_id,omitempty"` + ChatID string `json:"chat_id,omitempty"` + AssistantID string `json:"assistant_id,omitempty"` +} + +// Action the action +type Action struct { + Name string `json:"name,omitempty"` + Type string `json:"type"` + Payload interface{} `json:"payload,omitempty"` +} + +// New create a new message +func New() *Message { return &Message{Actions: []Action{}} } + +// NewString create a new message from string +func NewString(content string) (*Message, error) { + if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") { + var msg Message + if err := jsoniter.UnmarshalFromString(content, &msg); err != nil { + return nil, err + } + return &msg, nil + } + return &Message{Text: content}, nil +} + +// NewOpenAI create a new message from OpenAI response +func NewOpenAI(data []byte) *Message { + if data == nil || len(data) == 0 { + return nil + } + + msg := New() + text := string(data) + data = []byte(strings.TrimPrefix(text, "data: ")) + + switch { + case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"content":`): + var message openai.Message + if err := jsoniter.Unmarshal(data, &message); err != nil { + msg.Text = err.Error() + "\n" + string(data) + return msg + } + + if len(message.Choices) > 0 { + msg.Text = message.Choices[0].Delta.Content + } + + case strings.Contains(text, `[DONE]`): + msg.IsDone = true + + case strings.Contains(text, `"finish_reason":"stop"`): + msg.IsDone = true + + default: + str := strings.TrimPrefix(strings.Trim(string(data), "\""), "data: ") + msg.Type = "error" + msg.Text = str + } + + return msg +} + +// String returns the string representation +func (m *Message) String() string { + if m.Text != "" { + return m.Text + } + return "" +} + +// SetText set the text +func (m *Message) SetText(text string) *Message { + m.Text = text + if m.Data != nil { + if replaced := helper.Bind(text, m.Data); replaced != nil { + if replacedText, ok := replaced.(string); ok { + m.Text = replacedText + } + } + } + return m +} + +// Error set the error +func (m *Message) Error(message interface{}) *Message { + m.Type = "error" + switch v := message.(type) { + case error: + m.Text = v.Error() + case string: + m.Text = v + default: + m.Text = fmt.Sprintf("%v", message) + } + return m +} + +// Map set from map +func (m *Message) Map(msg map[string]interface{}) *Message { + if msg == nil { + return m + } + + if text, ok := msg["text"].(string); ok { + m.Text = text + } + if typ, ok := msg["type"].(string); ok { + m.Type = typ + } + if done, ok := msg["done"].(bool); ok { + m.IsDone = done + } + if actions, ok := msg["actions"].([]interface{}); ok { + for _, action := range actions { + if v, ok := action.(map[string]interface{}); ok { + action := Action{} + if name, ok := v["name"].(string); ok { + action.Name = name + } + if t, ok := v["type"].(string); ok { + action.Type = t + } + if payload, ok := v["payload"].(map[string]interface{}); ok { + action.Payload = payload + } + m.Actions = append(m.Actions, action) + } + } + } + if data, ok := msg["data"].(map[string]interface{}); ok { + m.Data = data + } + return m +} + +// Done set the done flag +func (m *Message) Done() *Message { + m.IsDone = true + return m +} + +// Action add an action +func (m *Message) Action(name string, t string, payload interface{}, next string) *Message { + if m.Data != nil { + payload = helper.Bind(payload, m.Data) + } + m.Actions = append(m.Actions, Action{ + Name: name, + Type: t, + Payload: payload, + }) + return m +} + +// Bind replace with data +func (m *Message) Bind(data map[string]interface{}) *Message { + if data == nil { + return m + } + m.Data = maps.Of(data).Dot() + return m +} + +// Write writes the message to response writer +func (m *Message) Write(w gin.ResponseWriter) bool { + defer func() { + if r := recover(); r != nil { + message := "Write Response Exception: (if client close the connection, it's normal) \n %s\n\n" + color.Red(message, r) + } + }() + + data, err := jsoniter.Marshal(m) + if err != nil { + log.Error("%s", err.Error()) + return false + } + + data = append([]byte("data: "), data...) + data = append(data, []byte("\n\n")...) + + if _, err := w.Write(data); err != nil { + color.Red("Write JSON Message Error: %s", err.Error()) + return false + } + w.Flush() + 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, "\"") + data := []byte(fmt.Sprintf(`{"text":"%s","type":"error"`, errMsg)) + if m.IsDone { + data = []byte(fmt.Sprintf(`{"text":"%s","type":"error","done":true`, errMsg)) + } + data = append([]byte("data: "), data...) + data = append(data, []byte("}\n\n")...) + + if _, err := w.Write(data); err != nil { + color.Red("Write JSON Message Error: %s", message) + } + w.Flush() +} + +// MarshalJSON implements json.Marshaler interface +func (m *Message) MarshalJSON() ([]byte, error) { + type Alias Message + return jsoniter.Marshal(&struct { + *Alias + }{ + Alias: (*Alias)(m), + }) +} + +// UnmarshalJSON implements json.Unmarshaler interface +func (m *Message) UnmarshalJSON(data []byte) error { + type Alias Message + aux := &struct { + *Alias + }{ + Alias: (*Alias)(m), + } + if err := jsoniter.Unmarshal(data, &aux); err != nil { + return err + } + return nil +} + +// MarshalJSON implements json.Marshaler interface +func (a *Action) MarshalJSON() ([]byte, error) { + type Alias Action + return jsoniter.Marshal(&struct { + *Alias + }{ + Alias: (*Alias)(a), + }) +} + +// UnmarshalJSON implements json.Unmarshaler interface +func (a *Action) UnmarshalJSON(data []byte) error { + type Alias Action + aux := &struct { + *Alias + }{ + Alias: (*Alias)(a), + } + if err := jsoniter.Unmarshal(data, &aux); err != nil { + return err + } + return nil +} diff --git a/neo/message/types.go b/neo/message/types.go deleted file mode 100644 index ed54cf31..00000000 --- a/neo/message/types.go +++ /dev/null @@ -1,27 +0,0 @@ -package message - -// Message the message -type Message struct { - Text string `json:"text,omitempty"` - Type string `json:"type,omitempty"` - Done bool `json:"done,omitempty"` - Confirm bool `json:"confirm,omitempty"` - Command *Command `json:"command,omitempty"` - Actions []Action `json:"actions,omitempty"` - Data map[string]interface{} `json:"-,omitempty"` -} - -// Action the action -type Action struct { - Name string `json:"name,omitempty"` - Type string `json:"type"` - Payload interface{} `json:"payload,omitempty"` - Next string `json:"next,omitempty"` -} - -// Command the command -type Command struct { - ID string `json:"id,omitempty"` - Name string `json:"name,omitempty"` - Reqeust string `json:"request,omitempty"` -} diff --git a/neo/neo.go b/neo/neo.go index 856082e6..0a034a67 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -124,7 +124,6 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy clientBreak := make(chan bool, 1) done := make(chan bool, 1) fail := make(chan error, 1) - content := []byte{} // Chat with AI in background @@ -142,26 +141,28 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy // Handle error if msg.Type == "error" { - fail <- fmt.Errorf("%s", msg.Message.Text) + fail <- fmt.Errorf("%s", msg.Text) return 0 // break } // Append content and send message content = msg.Append(content) - - // Only send real-time messages if not in silent mode - if !silent && msg.Message != nil && msg.Message.Text != "" { - message.New(). - Map(map[string]interface{}{ - "text": msg.Message.Text, - "done": msg.Message.Done, - }). - Write(c.Writer) + if !silent { + value := msg.String() + if value != "" { + message.New(). + Map(map[string]interface{}{ + "text": value, + "done": msg.IsDone, + }). + Write(c.Writer) + } } // Complete the stream - if msg.Message.Done { - if !silent && msg.Message.Text == "" { + if msg.IsDone { + value := msg.String() + if value == "" { msg.Write(c.Writer) } done <- true @@ -267,7 +268,6 @@ func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse, // chat chat with AI func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]interface{}, c *gin.Context) error { - if ast == nil { msg := message.New().Error("assistant is not initialized").Done() msg.Write(c.Writer) @@ -293,24 +293,26 @@ func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]inter // Handle error if msg.Type == "error" { - message.New().Error(msg.Message.Text).Done().Write(c.Writer) + value := msg.String() + message.New().Error(value).Done().Write(c.Writer) return 0 // break } // Append content and send message content = msg.Append(content) - if msg.Message != nil && msg.Message.Text != "" { + value := msg.String() + if value != "" { message.New(). Map(map[string]interface{}{ - "text": msg.Message.Text, - "done": msg.Message.Done, + "text": value, + "done": msg.IsDone, }). Write(c.Writer) } // Complete the stream - if msg.Message != nil && msg.Message.Done { - if msg.Message.Text == "" { + if msg.IsDone { + if value == "" { msg.Write(c.Writer) } done <- true @@ -544,9 +546,11 @@ func (neo *DSL) createConversation() error { // sendMessage sends a message to the client func (neo *DSL) sendMessage(w gin.ResponseWriter, data interface{}) error { - msg := message.New().Map(data.(map[string]interface{})) - if !msg.Write(w) { - return fmt.Errorf("failed to write message to stream") + 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 nil + return fmt.Errorf("invalid message data type") }