1package message
  2
  3import (
  4	"encoding/base64"
  5	"errors"
  6	"slices"
  7	"strings"
  8	"time"
  9
 10	"charm.land/fantasy"
 11	"charm.land/fantasy/providers/anthropic"
 12	"charm.land/fantasy/providers/google"
 13	"charm.land/fantasy/providers/openai"
 14	"github.com/charmbracelet/catwalk/pkg/catwalk"
 15)
 16
 17type MessageRole string
 18
 19const (
 20	Assistant MessageRole = "assistant"
 21	User      MessageRole = "user"
 22	System    MessageRole = "system"
 23	Tool      MessageRole = "tool"
 24)
 25
 26type FinishReason string
 27
 28const (
 29	FinishReasonEndTurn          FinishReason = "end_turn"
 30	FinishReasonMaxTokens        FinishReason = "max_tokens"
 31	FinishReasonToolUse          FinishReason = "tool_use"
 32	FinishReasonCanceled         FinishReason = "canceled"
 33	FinishReasonError            FinishReason = "error"
 34	FinishReasonPermissionDenied FinishReason = "permission_denied"
 35
 36	// Should never happen
 37	FinishReasonUnknown FinishReason = "unknown"
 38)
 39
 40type ContentPart interface {
 41	isPart()
 42}
 43
 44type ReasoningContent struct {
 45	Thinking         string                             `json:"thinking"`
 46	Signature        string                             `json:"signature"`
 47	ThoughtSignature string                             `json:"thought_signature"` // Used for google
 48	ResponsesData    *openai.ResponsesReasoningMetadata `json:"responses_data"`
 49	StartedAt        int64                              `json:"started_at,omitempty"`
 50	FinishedAt       int64                              `json:"finished_at,omitempty"`
 51}
 52
 53func (tc ReasoningContent) String() string {
 54	return tc.Thinking
 55}
 56func (ReasoningContent) isPart() {}
 57
 58type TextContent struct {
 59	Text string `json:"text"`
 60}
 61
 62func (tc TextContent) String() string {
 63	return tc.Text
 64}
 65
 66func (TextContent) isPart() {}
 67
 68type ImageURLContent struct {
 69	URL    string `json:"url"`
 70	Detail string `json:"detail,omitempty"`
 71}
 72
 73func (iuc ImageURLContent) String() string {
 74	return iuc.URL
 75}
 76
 77func (ImageURLContent) isPart() {}
 78
 79type BinaryContent struct {
 80	Path     string
 81	MIMEType string
 82	Data     []byte
 83}
 84
 85func (bc BinaryContent) String(p catwalk.InferenceProvider) string {
 86	base64Encoded := base64.StdEncoding.EncodeToString(bc.Data)
 87	if p == catwalk.InferenceProviderOpenAI {
 88		return "data:" + bc.MIMEType + ";base64," + base64Encoded
 89	}
 90	return base64Encoded
 91}
 92
 93func (BinaryContent) isPart() {}
 94
 95type ToolCall struct {
 96	ID               string `json:"id"`
 97	Name             string `json:"name"`
 98	Input            string `json:"input"`
 99	ProviderExecuted bool   `json:"provider_executed"`
100	Finished         bool   `json:"finished"`
101}
102
103func (ToolCall) isPart() {}
104
105type ToolResult struct {
106	ToolCallID string `json:"tool_call_id"`
107	Name       string `json:"name"`
108	Content    string `json:"content"`
109	Data       string `json:"data"`
110	MIMEType   string `json:"mime_type"`
111	Metadata   string `json:"metadata"`
112	IsError    bool   `json:"is_error"`
113}
114
115func (ToolResult) isPart() {}
116
117type Finish struct {
118	Reason  FinishReason `json:"reason"`
119	Time    int64        `json:"time"`
120	Message string       `json:"message,omitempty"`
121	Details string       `json:"details,omitempty"`
122}
123
124func (Finish) isPart() {}
125
126type Message struct {
127	ID               string
128	Role             MessageRole
129	SessionID        string
130	Parts            []ContentPart
131	Model            string
132	Provider         string
133	CreatedAt        int64
134	UpdatedAt        int64
135	IsSummaryMessage bool
136}
137
138func (m *Message) Content() TextContent {
139	for _, part := range m.Parts {
140		if c, ok := part.(TextContent); ok {
141			return c
142		}
143	}
144	return TextContent{}
145}
146
147func (m *Message) ReasoningContent() ReasoningContent {
148	for _, part := range m.Parts {
149		if c, ok := part.(ReasoningContent); ok {
150			return c
151		}
152	}
153	return ReasoningContent{}
154}
155
156func (m *Message) ImageURLContent() []ImageURLContent {
157	imageURLContents := make([]ImageURLContent, 0)
158	for _, part := range m.Parts {
159		if c, ok := part.(ImageURLContent); ok {
160			imageURLContents = append(imageURLContents, c)
161		}
162	}
163	return imageURLContents
164}
165
166func (m *Message) BinaryContent() []BinaryContent {
167	binaryContents := make([]BinaryContent, 0)
168	for _, part := range m.Parts {
169		if c, ok := part.(BinaryContent); ok {
170			binaryContents = append(binaryContents, c)
171		}
172	}
173	return binaryContents
174}
175
176func (m *Message) ToolCalls() []ToolCall {
177	toolCalls := make([]ToolCall, 0)
178	for _, part := range m.Parts {
179		if c, ok := part.(ToolCall); ok {
180			toolCalls = append(toolCalls, c)
181		}
182	}
183	return toolCalls
184}
185
186func (m *Message) ToolResults() []ToolResult {
187	toolResults := make([]ToolResult, 0)
188	for _, part := range m.Parts {
189		if c, ok := part.(ToolResult); ok {
190			toolResults = append(toolResults, c)
191		}
192	}
193	return toolResults
194}
195
196func (m *Message) IsFinished() bool {
197	for _, part := range m.Parts {
198		if _, ok := part.(Finish); ok {
199			return true
200		}
201	}
202	return false
203}
204
205func (m *Message) FinishPart() *Finish {
206	for _, part := range m.Parts {
207		if c, ok := part.(Finish); ok {
208			return &c
209		}
210	}
211	return nil
212}
213
214func (m *Message) FinishReason() FinishReason {
215	for _, part := range m.Parts {
216		if c, ok := part.(Finish); ok {
217			return c.Reason
218		}
219	}
220	return ""
221}
222
223func (m *Message) IsThinking() bool {
224	if m.ReasoningContent().Thinking != "" && m.Content().Text == "" && !m.IsFinished() {
225		return true
226	}
227	return false
228}
229
230func (m *Message) AppendContent(delta string) {
231	found := false
232	for i, part := range m.Parts {
233		if c, ok := part.(TextContent); ok {
234			m.Parts[i] = TextContent{Text: c.Text + delta}
235			found = true
236		}
237	}
238	if !found {
239		m.Parts = append(m.Parts, TextContent{Text: delta})
240	}
241}
242
243func (m *Message) AppendReasoningContent(delta string) {
244	found := false
245	for i, part := range m.Parts {
246		if c, ok := part.(ReasoningContent); ok {
247			m.Parts[i] = ReasoningContent{
248				Thinking:   c.Thinking + delta,
249				Signature:  c.Signature,
250				StartedAt:  c.StartedAt,
251				FinishedAt: c.FinishedAt,
252			}
253			found = true
254		}
255	}
256	if !found {
257		m.Parts = append(m.Parts, ReasoningContent{
258			Thinking:  delta,
259			StartedAt: time.Now().Unix(),
260		})
261	}
262}
263
264func (m *Message) AppendThoughtSignature(signature string) {
265	for i, part := range m.Parts {
266		if c, ok := part.(ReasoningContent); ok {
267			m.Parts[i] = ReasoningContent{
268				Thinking:         c.Thinking,
269				ThoughtSignature: c.ThoughtSignature + signature,
270				Signature:        c.Signature,
271				StartedAt:        c.StartedAt,
272				FinishedAt:       c.FinishedAt,
273			}
274			return
275		}
276	}
277	m.Parts = append(m.Parts, ReasoningContent{ThoughtSignature: signature})
278}
279
280func (m *Message) AppendReasoningSignature(signature string) {
281	for i, part := range m.Parts {
282		if c, ok := part.(ReasoningContent); ok {
283			m.Parts[i] = ReasoningContent{
284				Thinking:   c.Thinking,
285				Signature:  c.Signature + signature,
286				StartedAt:  c.StartedAt,
287				FinishedAt: c.FinishedAt,
288			}
289			return
290		}
291	}
292	m.Parts = append(m.Parts, ReasoningContent{Signature: signature})
293}
294
295func (m *Message) SetReasoningResponsesData(data *openai.ResponsesReasoningMetadata) {
296	for i, part := range m.Parts {
297		if c, ok := part.(ReasoningContent); ok {
298			m.Parts[i] = ReasoningContent{
299				Thinking:      c.Thinking,
300				ResponsesData: data,
301				StartedAt:     c.StartedAt,
302				FinishedAt:    c.FinishedAt,
303			}
304			return
305		}
306	}
307}
308
309func (m *Message) FinishThinking() {
310	for i, part := range m.Parts {
311		if c, ok := part.(ReasoningContent); ok {
312			if c.FinishedAt == 0 {
313				m.Parts[i] = ReasoningContent{
314					Thinking:   c.Thinking,
315					Signature:  c.Signature,
316					StartedAt:  c.StartedAt,
317					FinishedAt: time.Now().Unix(),
318				}
319			}
320			return
321		}
322	}
323}
324
325func (m *Message) ThinkingDuration() time.Duration {
326	reasoning := m.ReasoningContent()
327	if reasoning.StartedAt == 0 {
328		return 0
329	}
330
331	endTime := reasoning.FinishedAt
332	if endTime == 0 {
333		endTime = time.Now().Unix()
334	}
335
336	return time.Duration(endTime-reasoning.StartedAt) * time.Second
337}
338
339func (m *Message) FinishToolCall(toolCallID string) {
340	for i, part := range m.Parts {
341		if c, ok := part.(ToolCall); ok {
342			if c.ID == toolCallID {
343				m.Parts[i] = ToolCall{
344					ID:       c.ID,
345					Name:     c.Name,
346					Input:    c.Input,
347					Finished: true,
348				}
349				return
350			}
351		}
352	}
353}
354
355func (m *Message) AppendToolCallInput(toolCallID string, inputDelta string) {
356	for i, part := range m.Parts {
357		if c, ok := part.(ToolCall); ok {
358			if c.ID == toolCallID {
359				m.Parts[i] = ToolCall{
360					ID:       c.ID,
361					Name:     c.Name,
362					Input:    c.Input + inputDelta,
363					Finished: c.Finished,
364				}
365				return
366			}
367		}
368	}
369}
370
371func (m *Message) AddToolCall(tc ToolCall) {
372	for i, part := range m.Parts {
373		if c, ok := part.(ToolCall); ok {
374			if c.ID == tc.ID {
375				m.Parts[i] = tc
376				return
377			}
378		}
379	}
380	m.Parts = append(m.Parts, tc)
381}
382
383func (m *Message) SetToolCalls(tc []ToolCall) {
384	// remove any existing tool call part it could have multiple
385	parts := make([]ContentPart, 0)
386	for _, part := range m.Parts {
387		if _, ok := part.(ToolCall); ok {
388			continue
389		}
390		parts = append(parts, part)
391	}
392	m.Parts = parts
393	for _, toolCall := range tc {
394		m.Parts = append(m.Parts, toolCall)
395	}
396}
397
398func (m *Message) AddToolResult(tr ToolResult) {
399	m.Parts = append(m.Parts, tr)
400}
401
402func (m *Message) SetToolResults(tr []ToolResult) {
403	for _, toolResult := range tr {
404		m.Parts = append(m.Parts, toolResult)
405	}
406}
407
408func (m *Message) AddFinish(reason FinishReason, message, details string) {
409	// remove any existing finish part
410	for i, part := range m.Parts {
411		if _, ok := part.(Finish); ok {
412			m.Parts = slices.Delete(m.Parts, i, i+1)
413			break
414		}
415	}
416	m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now().Unix(), Message: message, Details: details})
417}
418
419func (m *Message) AddImageURL(url, detail string) {
420	m.Parts = append(m.Parts, ImageURLContent{URL: url, Detail: detail})
421}
422
423func (m *Message) AddBinary(mimeType string, data []byte) {
424	m.Parts = append(m.Parts, BinaryContent{MIMEType: mimeType, Data: data})
425}
426
427func (m *Message) ToAIMessage() []fantasy.Message {
428	var messages []fantasy.Message
429	switch m.Role {
430	case User:
431		var parts []fantasy.MessagePart
432		text := strings.TrimSpace(m.Content().Text)
433		if text != "" {
434			parts = append(parts, fantasy.TextPart{Text: text})
435		}
436		for _, content := range m.BinaryContent() {
437			parts = append(parts, fantasy.FilePart{
438				Filename:  content.Path,
439				Data:      content.Data,
440				MediaType: content.MIMEType,
441			})
442		}
443		messages = append(messages, fantasy.Message{
444			Role:    fantasy.MessageRoleUser,
445			Content: parts,
446		})
447	case Assistant:
448		var parts []fantasy.MessagePart
449		text := strings.TrimSpace(m.Content().Text)
450		if text != "" {
451			parts = append(parts, fantasy.TextPart{Text: text})
452		}
453		reasoning := m.ReasoningContent()
454		if reasoning.Thinking != "" {
455			reasoningPart := fantasy.ReasoningPart{Text: reasoning.Thinking, ProviderOptions: fantasy.ProviderOptions{}}
456			if reasoning.Signature != "" {
457				reasoningPart.ProviderOptions[anthropic.Name] = &anthropic.ReasoningOptionMetadata{
458					Signature: reasoning.Signature,
459				}
460			}
461			if reasoning.ResponsesData != nil {
462				reasoningPart.ProviderOptions[openai.Name] = reasoning.ResponsesData
463			}
464			if reasoning.ThoughtSignature != "" {
465				reasoningPart.ProviderOptions[google.Name] = &google.ReasoningMetadata{
466					Signature: reasoning.ThoughtSignature,
467				}
468			}
469			parts = append(parts, reasoningPart)
470		}
471		for _, call := range m.ToolCalls() {
472			parts = append(parts, fantasy.ToolCallPart{
473				ToolCallID:       call.ID,
474				ToolName:         call.Name,
475				Input:            call.Input,
476				ProviderExecuted: call.ProviderExecuted,
477			})
478		}
479		messages = append(messages, fantasy.Message{
480			Role:    fantasy.MessageRoleAssistant,
481			Content: parts,
482		})
483	case Tool:
484		var parts []fantasy.MessagePart
485		for _, result := range m.ToolResults() {
486			var content fantasy.ToolResultOutputContent
487			if result.IsError {
488				content = fantasy.ToolResultOutputContentError{
489					Error: errors.New(result.Content),
490				}
491			} else if result.Data != "" {
492				content = fantasy.ToolResultOutputContentMedia{
493					Data:      result.Data,
494					MediaType: result.MIMEType,
495				}
496			} else {
497				content = fantasy.ToolResultOutputContentText{
498					Text: result.Content,
499				}
500			}
501			parts = append(parts, fantasy.ToolResultPart{
502				ToolCallID: result.ToolCallID,
503				Output:     content,
504			})
505		}
506		messages = append(messages, fantasy.Message{
507			Role:    fantasy.MessageRoleTool,
508			Content: parts,
509		})
510	}
511	return messages
512}