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