lmcli/pkg/cli/conversation.go

package cli

import (
	"fmt"
	"strings"
)

type MessageRole string

const (
	MessageRoleUser      MessageRole = "user"
	MessageRoleAssistant MessageRole = "assistant"
	MessageRoleSystem    MessageRole = "system"
)

// FriendlyRole returns a human friendly signifier for the message's role.
func (m *Message) FriendlyRole() string {
	var friendlyRole string
	switch m.Role {
	case MessageRoleUser:
		friendlyRole = "You"
	case MessageRoleSystem:
		friendlyRole = "System"
	case MessageRoleAssistant:
		friendlyRole = "Assistant"
	default:
		friendlyRole = string(m.Role)
	}
	return friendlyRole
}

func (c *Conversation) GenerateTitle() (string, error) {
	messages, err := store.Messages(c)
	if err != nil {
		return "", err
	}

	const header = "Generate a concise 4-5 word title for the conversation below."
	prompt := fmt.Sprintf("%s\n\n---\n\n%s", header, formatForExternalPrompting(messages, false))

	generateRequest := []Message{
		{
			Role:            MessageRoleUser,
			OriginalContent: prompt,
		},
	}

	model := "gpt-3.5-turbo" // use cheap model to generate title
	response, err := CreateChatCompletion(model, generateRequest, 25, nil)
	if err != nil {
		return "", err
	}

	return response, nil
}

func formatForExternalPrompting(messages []Message, system bool) string {
	sb := strings.Builder{}
	for _, message := range messages {
		if message.Role == MessageRoleSystem && !system {
			continue
		}
		sb.WriteString(fmt.Sprintf("<%s>\n", message.FriendlyRole()))
		sb.WriteString(fmt.Sprintf("\"\"\"\n%s\n\"\"\"\n\n", message.OriginalContent))
	}
	return sb.String()
}