message.go

package blades

import (
	"fmt"
	"maps"
	"slices"
	"strings"

	"github.com/google/uuid"
)

// Role indicates the author of a message in a conversation.
type Role string

const (
	// RoleUser is an end user.
	RoleUser Role = "user"
	// RoleSystem provides system-level instructions.
	RoleSystem Role = "system"
	// RoleAssistant is the model output.
	RoleAssistant Role = "assistant"
	// RoleTool indicates a message generated by a tool.
	RoleTool Role = "tool"
)

// Status indicates the state of a message.
type Status string

const (
	// StatusInProgress indicates the message is being generated.
	StatusInProgress Status = "in_progress"
	// StatusIncomplete indicates the message is partially generated.
	StatusIncomplete Status = "incomplete"
	// StatusCompleted indicates the message is fully generated.
	StatusCompleted Status = "completed"
)

// TextPart is plain text content.
type TextPart struct {
	Text string `json:"text"`
}

// FilePart is a reference to a file by its URI.
type FilePart struct {
	Name     string   `json:"name"`
	URI      string   `json:"uri"`
	MIMEType MIMEType `json:"mimeType"`
}

// DataPart is a file represented by its byte content.
type DataPart struct {
	Name     string   `json:"name"`
	Bytes    []byte   `json:"bytes"`
	MIMEType MIMEType `json:"mimeType"`
}

// ToolPart is a tool call, containing its request, response, and completion state.
type ToolPart struct {
	ID        string `json:"id"`
	Name      string `json:"name"`
	Request   string `json:"arguments"`
	Response  string `json:"result,omitempty"`
	Completed bool   `json:"completed,omitempty"`
}

// NewToolPart creates a tool call part that has not completed yet.
func NewToolPart(id, name, request string) ToolPart {
	return ToolPart{
		ID:      id,
		Name:    name,
		Request: request,
	}
}

// Part is a part of a message, which can be text or a file.
type Part interface {
	isPart()
}

func (TextPart) isPart() {}
func (FilePart) isPart() {}
func (DataPart) isPart() {}
func (ToolPart) isPart() {}

// TokenUsage tracks token consumption for a message.
type TokenUsage struct {
	InputTokens  int64 `json:"inputTokens"`
	OutputTokens int64 `json:"outputTokens"`
	TotalTokens  int64 `json:"totalTokens"`
}

// Message represents a single message in a conversation.
type Message struct {
	ID           string         `json:"id"`
	Role         Role           `json:"role"`
	Parts        []Part         `json:"parts"`
	Author       string         `json:"author"`
	InvocationID string         `json:"invocationId,omitempty"`
	Status       Status         `json:"status"`
	FinishReason string         `json:"finishReason,omitempty"`
	TokenUsage   TokenUsage     `json:"tokenUsage,omitempty"`
	Actions      map[string]any `json:"actions,omitempty"`
	Metadata     map[string]any `json:"metadata,omitempty"`
}

// Text returns the first text part of the message, or an empty string if none exists.
func (m *Message) Text() string {
	var buf strings.Builder
	for _, part := range m.Parts {
		switch v := any(part).(type) {
		case TextPart:
			buf.WriteString(v.Text)
			buf.WriteByte('\n')
		}
	}
	return strings.TrimSuffix(buf.String(), "\n")
}

// File returns the first file part of the message, or nil if none exists.
func (m *Message) File() *FilePart {
	for _, part := range m.Parts {
		if file, ok := part.(FilePart); ok {
			return &file
		}
	}
	return nil
}

// Data returns the first data part of the message, or nil if none exists.
func (m *Message) Data() *DataPart {
	for _, part := range m.Parts {
		if data, ok := part.(DataPart); ok {
			return &data
		}
	}
	return nil
}

// Clone creates a deep copy of the message.
func (m *Message) Clone() *Message {
	if m == nil {
		return nil
	}
	c := *m
	c.Parts = slices.Clone(m.Parts)
	if m.Actions != nil {
		c.Actions = maps.Clone(m.Actions)
	}
	if m.Metadata != nil {
		c.Metadata = maps.Clone(m.Metadata)
	}
	return &c
}

func (m *Message) String() string {
	var buf strings.Builder
	for _, part := range m.Parts {
		switch v := part.(type) {
		case TextPart:
			buf.WriteString("[Text: " + v.Text + "]")
		case FilePart:
			buf.WriteString("[File: " + v.Name + " (" + string(v.MIMEType) + ")]")
		case DataPart:
			buf.WriteString("[Data: " + v.Name + " (" + string(v.MIMEType) + "), " + fmt.Sprintf("%d bytes", len(v.Bytes)) + "]")
		case ToolPart:
			buf.WriteString("[Tool: " + v.Name + " (Request: " + v.Request + ", Response: " + v.Response + ")]")
		}
	}
	return buf.String()
}

// UserMessage creates a user-authored message from parts.
func UserMessage(parts ...any) *Message {
	return &Message{ID: NewMessageID(), Role: RoleUser, Author: "user", Parts: NewMessageParts(parts...)}
}

// SystemMessage creates a system-authored message from parts.
func SystemMessage(parts ...any) *Message {
	return &Message{ID: NewMessageID(), Role: RoleSystem, Parts: NewMessageParts(parts...)}
}

// AssistantMessage creates an assistant-authored message from parts.
func AssistantMessage(parts ...any) *Message {
	return &Message{ID: NewMessageID(), Role: RoleAssistant, Parts: NewMessageParts(parts...)}
}

// NewAssistantMessage creates a new assistant message with the given status.
func NewAssistantMessage(status Status) *Message {
	return &Message{
		ID:       NewMessageID(),
		Role:     RoleAssistant,
		Status:   status,
		Actions:  make(map[string]any),
		Metadata: make(map[string]any),
	}
}

// NewMessageID generates a new random message identifier.
func NewMessageID() string {
	return uuid.NewString()
}

// NewMessageParts converts a heterogeneous list of content inputs into model parts.
// Accepts raw string, Text, FileURI, and FileBytes.
func NewMessageParts(inputs ...any) []Part {
	parts := make([]Part, 0, len(inputs))
	for _, input := range inputs {
		switch v := any(input).(type) {
		case string:
			parts = append(parts, TextPart{v})
		case TextPart:
			parts = append(parts, v)
		case FilePart:
			parts = append(parts, v)
		case DataPart:
			parts = append(parts, v)
		case ToolPart:
			parts = append(parts, v)
		}
	}
	return parts
}

// MergeParts merges the Parts of two messages, mutating and returning base.
// If base is nil, extra is returned. If extra is nil, base is returned.
func MergeParts(base, extra *Message) *Message {
	if base == nil {
		return extra
	}
	if extra == nil {
		return base
	}
	base.Parts = append(base.Parts, extra.Parts...)
	return base
}

// MergeActions merges two action maps, with values from extra overriding those in base.
func MergeActions(base, extra map[string]any) map[string]any {
	actions := make(map[string]any, len(base)+len(extra))
	for k, v := range base {
		actions[k] = v
	}
	for k, v := range extra {
		actions[k] = v
	}
	return actions
}

// AppendMessages appends extra messages to base, replacing any messages in base that have matching IDs in extra.
func AppendMessages(base []*Message, extra ...*Message) []*Message {
	var (
		sets     = make(map[string]struct{}, len(extra))
		filtered = make([]*Message, 0, len(base))
	)
	for _, m := range extra {
		if m.ID == "" {
			continue
		}
		sets[m.ID] = struct{}{}
	}
	for _, m := range base {
		if _, exists := sets[m.ID]; !exists {
			filtered = append(filtered, m)
		}
	}
	return append(filtered, extra...)
}