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

@@ -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
+}

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
+}

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
+}

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

neo/message/json.go

@@ -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()
-}

neo/message/message.go

@@ -1,6 +1,296 @@
 package message
 
-// makeMessage create a new message
+import (
-func makeMessage() *Message {
+	"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
+}

neo/message/types.go

@@ -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"`
-}

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)
-
+				if !silent {
-				// Only send real-time messages if not in silent mode
+					value := msg.String()
-				if !silent && msg.Message != nil && msg.Message.Text != "" {
+					if value != "" {
-					message.New().
+						message.New().
-						Map(map[string]interface{}{
+							Map(map[string]interface{}{
-							"text": msg.Message.Text,
+								"text": value,
-							"done": msg.Message.Done,
+								"done": msg.IsDone,
-						}).
+							}).
-						Write(c.Writer)
+							Write(c.Writer)
+					}
 				}
 
 				// Complete the stream
-				if msg.Message.Done {
+				if msg.IsDone {
-					if !silent && msg.Message.Text == "" {
+					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,
+							"text": value,
-							"done": msg.Message.Done,
+							"done": msg.IsDone,
 						}).
 						Write(c.Writer)
 				}
 
 				// Complete the stream
-				if msg.Message != nil && msg.Message.Done {
+				if msg.IsDone {
-					if msg.Message.Text == "" {
+					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, ok := data.(map[string]interface{}); ok {
-	if !msg.Write(w) {
+		if !message.New().Map(msg).Write(w) {
-		return fmt.Errorf("failed to write message to stream")
+			return fmt.Errorf("failed to write message to stream")
+		}
+		return nil
 	}
-	return nil
+	return fmt.Errorf("invalid message data type")
 }