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 }