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