content.go

  1package message
  2
  3import (
  4	"encoding/base64"
  5	"time"
  6)
  7
  8type MessageRole string
  9
 10const (
 11	Assistant MessageRole = "assistant"
 12	User      MessageRole = "user"
 13	System    MessageRole = "system"
 14	Tool      MessageRole = "tool"
 15)
 16
 17type ContentPart interface {
 18	isPart()
 19}
 20
 21type ReasoningContent struct {
 22	Thinking string `json:"thinking"`
 23}
 24
 25func (tc ReasoningContent) String() string {
 26	return tc.Thinking
 27}
 28func (ReasoningContent) isPart() {}
 29
 30type TextContent struct {
 31	Text string `json:"text"`
 32}
 33
 34func (tc TextContent) String() string {
 35	return tc.Text
 36}
 37
 38func (TextContent) isPart() {}
 39
 40type ImageURLContent struct {
 41	URL    string `json:"url"`
 42	Detail string `json:"detail,omitempty"`
 43}
 44
 45func (iuc ImageURLContent) String() string {
 46	return iuc.URL
 47}
 48
 49func (ImageURLContent) isPart() {}
 50
 51type BinaryContent struct {
 52	MIMEType string
 53	Data     []byte
 54}
 55
 56func (bc BinaryContent) String() string {
 57	base64Encoded := base64.StdEncoding.EncodeToString(bc.Data)
 58	return "data:" + bc.MIMEType + ";base64," + base64Encoded
 59}
 60
 61func (BinaryContent) isPart() {}
 62
 63type ToolCall struct {
 64	ID       string `json:"id"`
 65	Name     string `json:"name"`
 66	Input    string `json:"input"`
 67	Type     string `json:"type"`
 68	Metadata any    `json:"metadata"`
 69	Finished bool   `json:"finished"`
 70}
 71
 72func (ToolCall) isPart() {}
 73
 74type ToolResult struct {
 75	ToolCallID string `json:"tool_call_id"`
 76	Name       string `json:"name"`
 77	Content    string `json:"content"`
 78	IsError    bool   `json:"is_error"`
 79}
 80
 81func (ToolResult) isPart() {}
 82
 83type Finish struct {
 84	Reason string `json:"reason"`
 85	Time   int64  `json:"time"`
 86}
 87
 88func (Finish) isPart() {}
 89
 90type Message struct {
 91	ID        string
 92	Role      MessageRole
 93	SessionID string
 94	Parts     []ContentPart
 95
 96	CreatedAt int64
 97	UpdatedAt int64
 98}
 99
100func (m *Message) Content() TextContent {
101	for _, part := range m.Parts {
102		if c, ok := part.(TextContent); ok {
103			return c
104		}
105	}
106	return TextContent{}
107}
108
109func (m *Message) ReasoningContent() ReasoningContent {
110	for _, part := range m.Parts {
111		if c, ok := part.(ReasoningContent); ok {
112			return c
113		}
114	}
115	return ReasoningContent{}
116}
117
118func (m *Message) ImageURLContent() []ImageURLContent {
119	imageURLContents := make([]ImageURLContent, 0)
120	for _, part := range m.Parts {
121		if c, ok := part.(ImageURLContent); ok {
122			imageURLContents = append(imageURLContents, c)
123		}
124	}
125	return imageURLContents
126}
127
128func (m *Message) BinaryContent() []BinaryContent {
129	binaryContents := make([]BinaryContent, 0)
130	for _, part := range m.Parts {
131		if c, ok := part.(BinaryContent); ok {
132			binaryContents = append(binaryContents, c)
133		}
134	}
135	return binaryContents
136}
137
138func (m *Message) ToolCalls() []ToolCall {
139	toolCalls := make([]ToolCall, 0)
140	for _, part := range m.Parts {
141		if c, ok := part.(ToolCall); ok {
142			toolCalls = append(toolCalls, c)
143		}
144	}
145	return toolCalls
146}
147
148func (m *Message) ToolResults() []ToolResult {
149	toolResults := make([]ToolResult, 0)
150	for _, part := range m.Parts {
151		if c, ok := part.(ToolResult); ok {
152			toolResults = append(toolResults, c)
153		}
154	}
155	return toolResults
156}
157
158func (m *Message) IsFinished() bool {
159	for _, part := range m.Parts {
160		if _, ok := part.(Finish); ok {
161			return true
162		}
163	}
164	return false
165}
166
167func (m *Message) FinishPart() *Finish {
168	for _, part := range m.Parts {
169		if c, ok := part.(Finish); ok {
170			return &c
171		}
172	}
173	return nil
174}
175
176func (m *Message) FinishReason() string {
177	for _, part := range m.Parts {
178		if c, ok := part.(Finish); ok {
179			return c.Reason
180		}
181	}
182	return ""
183}
184
185func (m *Message) IsThinking() bool {
186	if m.ReasoningContent().Thinking != "" && m.Content().Text == "" && !m.IsFinished() {
187		return true
188	}
189	return false
190}
191
192func (m *Message) AppendContent(delta string) {
193	found := false
194	for i, part := range m.Parts {
195		if c, ok := part.(TextContent); ok {
196			m.Parts[i] = TextContent{Text: c.Text + delta}
197			found = true
198		}
199	}
200	if !found {
201		m.Parts = append(m.Parts, TextContent{Text: delta})
202	}
203}
204
205func (m *Message) AppendReasoningContent(delta string) {
206	found := false
207	for i, part := range m.Parts {
208		if c, ok := part.(ReasoningContent); ok {
209			m.Parts[i] = ReasoningContent{Thinking: c.Thinking + delta}
210			found = true
211		}
212	}
213	if !found {
214		m.Parts = append(m.Parts, ReasoningContent{Thinking: delta})
215	}
216}
217
218func (m *Message) AddToolCall(tc ToolCall) {
219	for i, part := range m.Parts {
220		if c, ok := part.(ToolCall); ok {
221			if c.ID == tc.ID {
222				m.Parts[i] = tc
223				return
224			}
225		}
226	}
227	m.Parts = append(m.Parts, tc)
228}
229
230func (m *Message) SetToolCalls(tc []ToolCall) {
231	for _, toolCall := range tc {
232		m.Parts = append(m.Parts, toolCall)
233	}
234}
235
236func (m *Message) AddToolResult(tr ToolResult) {
237	m.Parts = append(m.Parts, tr)
238}
239
240func (m *Message) SetToolResults(tr []ToolResult) {
241	for _, toolResult := range tr {
242		m.Parts = append(m.Parts, toolResult)
243	}
244}
245
246func (m *Message) AddFinish(reason string) {
247	m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now().Unix()})
248}
249
250func (m *Message) AddImageURL(url, detail string) {
251	m.Parts = append(m.Parts, ImageURLContent{URL: url, Detail: detail})
252}
253
254func (m *Message) AddBinary(mimeType string, data []byte) {
255	m.Parts = append(m.Parts, BinaryContent{MIMEType: mimeType, Data: data})
256}