yao/neo/message/message.go

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"
)

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