lmcli/pkg/api/message.go

package api

import (
	"time"
)

type MessageRole string

const (
	MessageRoleSystem     MessageRole = "system"
	MessageRoleUser       MessageRole = "user"
	MessageRoleAssistant  MessageRole = "assistant"
	MessageRoleToolCall   MessageRole = "tool_call"
	MessageRoleToolResult MessageRole = "tool_result"
)

type Message struct {
	ID             uint          `gorm:"primaryKey"`
	ConversationID *uint         `gorm:"index"`
	Conversation   *Conversation `gorm:"foreignKey:ConversationID"`
	Content        string
	Role           MessageRole
	CreatedAt      time.Time
	ToolCalls      ToolCalls   // a json array of tool calls (from the model)
	ToolResults    ToolResults // a json array of tool results
	ParentID       *uint
	Parent         *Message  `gorm:"foreignKey:ParentID"`
	Replies        []Message `gorm:"foreignKey:ParentID"`

	SelectedReplyID *uint
	SelectedReply   *Message `gorm:"foreignKey:SelectedReplyID"`
}

func ApplySystemPrompt(m []Message, system string, force bool) []Message {
	if len(m) > 0 && m[0].Role == MessageRoleSystem {
		if force {
			m[0].Content = system
		}
		return m
	} else {
		return append([]Message{{
			Role:    MessageRoleSystem,
			Content: system,
		}}, m...)
	}
}

func (m *MessageRole) IsAssistant() bool {
	switch *m {
	case MessageRoleAssistant, MessageRoleToolCall:
		return true
	}
	return false
}

// FriendlyRole returns a human friendly signifier for the message's role.
func (m MessageRole) FriendlyRole() string {
	switch m {
	case MessageRoleUser:
		return "You"
	case MessageRoleSystem:
		return "System"
	case MessageRoleAssistant:
		return "Assistant"
	case MessageRoleToolCall:
		return "Tool Call"
	case MessageRoleToolResult:
		return "Tool Result"
	default:
		return string(m)
	}
}