lmcli/pkg/api/api.go

package api

import (
	"encoding/json"
	"fmt"
)

type MessageRole string

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

type Message struct {
	Content     string // TODO: support multi-part messages
	Role        MessageRole
	ToolCalls   []ToolCall
	ToolResults []ToolResult
}

type ToolSpec struct {
	Name        string
	Description string
	Parameters  []ToolParameter
	Impl        func(*ToolSpec, map[string]interface{}) (string, error)
}

type ToolParameter struct {
	Name        string   `json:"name"`
	Type        string   `json:"type"` // "string", "integer", "boolean"
	Required    bool     `json:"required"`
	Description string   `json:"description"`
	Enum        []string `json:"enum,omitempty"`
}

type ToolCall struct {
	ID         string                 `json:"id" yaml:"-"`
	Name       string                 `json:"name" yaml:"tool"`
	Parameters map[string]interface{} `json:"parameters" yaml:"parameters"`
}

type ToolResult struct {
	ToolCallID string `json:"toolCallID" yaml:"-"`
	ToolName   string `json:"toolName,omitempty" yaml:"tool"`
	Result     string `json:"result,omitempty" yaml:"result"`
}

func NewMessageWithAssistant(content string) *Message {
	return &Message{
		Role:    MessageRoleAssistant,
		Content: content,
	}
}

func NewMessageWithToolCalls(content string, toolCalls []ToolCall) *Message {
	return &Message{
		Role:      MessageRoleToolCall,
		Content:   content,
		ToolCalls: toolCalls,
	}
}

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

func (m MessageRole) IsUser() bool {
	switch m {
	case MessageRoleUser, MessageRoleToolResult:
		return true
	}
	return false
}

func (m MessageRole) IsSystem() bool {
	switch m {
	case MessageRoleSystem:
		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)
	}
}

// TODO: remove this
type CallResult struct {
	Message string `json:"message"`
	Result  any    `json:"result,omitempty"`
}

func (r CallResult) ToJson() (string, error) {
	if r.Message == "" {
		// When message not supplied, assume success
		r.Message = "success"
	}

	jsonBytes, err := json.Marshal(r)
	if err != nil {
		return "", fmt.Errorf("Could not marshal CallResult to JSON: %v\n", err)
	}
	return string(jsonBytes), nil
}