434 lines
12 KiB
Go
434 lines
12 KiB
Go
package lmcli
|
|
|
|
import (
|
|
"database/sql"
|
|
"errors"
|
|
"fmt"
|
|
"slices"
|
|
"strings"
|
|
"time"
|
|
|
|
"git.mlow.ca/mlow/lmcli/pkg/api"
|
|
sqids "github.com/sqids/sqids-go"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type ConversationStore interface {
|
|
ConversationByShortName(shortName string) (*api.Conversation, error)
|
|
ConversationShortNameCompletions(search string) []string
|
|
RootMessages(conversationID uint) ([]api.Message, error)
|
|
LatestConversationMessages() ([]api.Message, error)
|
|
|
|
StartConversation(messages ...api.Message) (*api.Conversation, []api.Message, error)
|
|
UpdateConversation(conversation *api.Conversation) error
|
|
DeleteConversation(conversation *api.Conversation) error
|
|
CloneConversation(toClone api.Conversation) (*api.Conversation, uint, error)
|
|
|
|
MessageByID(messageID uint) (*api.Message, error)
|
|
MessageReplies(messageID uint) ([]api.Message, error)
|
|
|
|
UpdateMessage(message *api.Message) error
|
|
DeleteMessage(message *api.Message, prune bool) error
|
|
CloneBranch(toClone api.Message) (*api.Message, uint, error)
|
|
Reply(to *api.Message, messages ...api.Message) ([]api.Message, error)
|
|
|
|
PathToRoot(message *api.Message) ([]api.Message, error)
|
|
PathToLeaf(message *api.Message) ([]api.Message, error)
|
|
}
|
|
|
|
type SQLStore struct {
|
|
db *gorm.DB
|
|
sqids *sqids.Sqids
|
|
}
|
|
|
|
func NewSQLStore(db *gorm.DB) (*SQLStore, error) {
|
|
models := []any{
|
|
&api.Conversation{},
|
|
&api.Message{},
|
|
}
|
|
|
|
for _, x := range models {
|
|
err := db.AutoMigrate(x)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Could not perform database migrations: %v", err)
|
|
}
|
|
}
|
|
|
|
_sqids, _ := sqids.New(sqids.Options{MinLength: 4})
|
|
return &SQLStore{db, _sqids}, nil
|
|
}
|
|
|
|
func (s *SQLStore) createConversation() (*api.Conversation, error) {
|
|
// Create the new conversation
|
|
c := &api.Conversation{}
|
|
err := s.db.Save(c).Error
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// Generate and save its "short name"
|
|
shortName, _ := s.sqids.Encode([]uint64{uint64(c.ID)})
|
|
c.ShortName = sql.NullString{String: shortName, Valid: true}
|
|
err = s.db.Updates(c).Error
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return c, nil
|
|
}
|
|
|
|
func (s *SQLStore) UpdateConversation(c *api.Conversation) error {
|
|
if c == nil || c.ID == 0 {
|
|
return fmt.Errorf("Conversation is nil or invalid (missing ID)")
|
|
}
|
|
return s.db.Updates(c).Error
|
|
}
|
|
|
|
func (s *SQLStore) DeleteConversation(c *api.Conversation) error {
|
|
// Delete messages first
|
|
err := s.db.Where("conversation_id = ?", c.ID).Delete(&api.Message{}).Error
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return s.db.Delete(c).Error
|
|
}
|
|
|
|
func (s *SQLStore) DeleteMessage(message *api.Message, prune bool) error {
|
|
panic("Not yet implemented")
|
|
//return s.db.Delete(&message).Error
|
|
}
|
|
|
|
func (s *SQLStore) UpdateMessage(m *api.Message) error {
|
|
if m == nil || m.ID == 0 {
|
|
return fmt.Errorf("Message is nil or invalid (missing ID)")
|
|
}
|
|
return s.db.Updates(m).Error
|
|
}
|
|
|
|
func (s *SQLStore) ConversationShortNameCompletions(shortName string) []string {
|
|
var conversations []api.Conversation
|
|
// ignore error for completions
|
|
s.db.Find(&conversations)
|
|
completions := make([]string, 0, len(conversations))
|
|
for _, conversation := range conversations {
|
|
if shortName == "" || strings.HasPrefix(conversation.ShortName.String, shortName) {
|
|
completions = append(completions, fmt.Sprintf("%s\t%s", conversation.ShortName.String, conversation.Title))
|
|
}
|
|
}
|
|
return completions
|
|
}
|
|
|
|
func (s *SQLStore) ConversationByShortName(shortName string) (*api.Conversation, error) {
|
|
if shortName == "" {
|
|
return nil, errors.New("shortName is empty")
|
|
}
|
|
var conversation api.Conversation
|
|
err := s.db.Preload("SelectedRoot").Where("short_name = ?", shortName).Find(&conversation).Error
|
|
return &conversation, err
|
|
}
|
|
|
|
func (s *SQLStore) RootMessages(conversationID uint) ([]api.Message, error) {
|
|
var rootMessages []api.Message
|
|
err := s.db.Where("conversation_id = ? AND parent_id IS NULL", conversationID).Find(&rootMessages).Error
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return rootMessages, nil
|
|
}
|
|
|
|
func (s *SQLStore) MessageByID(messageID uint) (*api.Message, error) {
|
|
var message api.Message
|
|
err := s.db.Preload("Parent").Preload("Replies").Preload("SelectedReply").Where("id = ?", messageID).Find(&message).Error
|
|
return &message, err
|
|
}
|
|
|
|
func (s *SQLStore) MessageReplies(messageID uint) ([]api.Message, error) {
|
|
var replies []api.Message
|
|
err := s.db.Where("parent_id = ?", messageID).Find(&replies).Error
|
|
return replies, err
|
|
}
|
|
|
|
// StartConversation starts a new conversation with the provided messages
|
|
func (s *SQLStore) StartConversation(messages ...api.Message) (*api.Conversation, []api.Message, error) {
|
|
if len(messages) == 0 {
|
|
return nil, nil, fmt.Errorf("Must provide at least 1 message")
|
|
}
|
|
|
|
// Create new conversation
|
|
conversation, err := s.createConversation()
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
// Create first message
|
|
messages[0].Conversation = conversation
|
|
err = s.db.Create(&messages[0]).Error
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
// Update conversation's selected root message
|
|
conversation.SelectedRoot = &messages[0]
|
|
err = s.UpdateConversation(conversation)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
// Add additional replies to conversation
|
|
if len(messages) > 1 {
|
|
newMessages, err := s.Reply(&messages[0], messages[1:]...)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
messages = append([]api.Message{messages[0]}, newMessages...)
|
|
}
|
|
return conversation, messages, nil
|
|
}
|
|
|
|
// CloneConversation clones the given conversation and all of its root meesages
|
|
func (s *SQLStore) CloneConversation(toClone api.Conversation) (*api.Conversation, uint, error) {
|
|
rootMessages, err := s.RootMessages(toClone.ID)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
|
|
clone, err := s.createConversation()
|
|
if err != nil {
|
|
return nil, 0, fmt.Errorf("Could not create clone: %s", err)
|
|
}
|
|
clone.Title = toClone.Title + " - Clone"
|
|
|
|
var errors []error
|
|
var messageCnt uint = 0
|
|
for _, root := range rootMessages {
|
|
messageCnt++
|
|
newRoot := root
|
|
newRoot.ConversationID = &clone.ID
|
|
|
|
cloned, count, err := s.CloneBranch(newRoot)
|
|
if err != nil {
|
|
errors = append(errors, err)
|
|
continue
|
|
}
|
|
messageCnt += count
|
|
|
|
if root.ID == *toClone.SelectedRootID {
|
|
clone.SelectedRootID = &cloned.ID
|
|
if err := s.UpdateConversation(clone); err != nil {
|
|
errors = append(errors, fmt.Errorf("Could not set selected root on clone: %v", err))
|
|
}
|
|
}
|
|
}
|
|
|
|
if len(errors) > 0 {
|
|
return nil, 0, fmt.Errorf("Messages failed to be cloned: %v", errors)
|
|
}
|
|
|
|
return clone, messageCnt, nil
|
|
}
|
|
|
|
// Reply to a message with a series of messages (each following the next)
|
|
func (s *SQLStore) Reply(to *api.Message, messages ...api.Message) ([]api.Message, error) {
|
|
var savedMessages []api.Message
|
|
|
|
err := s.db.Transaction(func(tx *gorm.DB) error {
|
|
currentParent := to
|
|
for i := range messages {
|
|
parent := currentParent
|
|
message := messages[i]
|
|
message.Parent = parent
|
|
message.Conversation = parent.Conversation
|
|
message.ID = 0
|
|
message.CreatedAt = time.Time{}
|
|
|
|
if err := tx.Create(&message).Error; err != nil {
|
|
return err
|
|
}
|
|
|
|
// update parent selected reply
|
|
parent.Replies = append(parent.Replies, message)
|
|
parent.SelectedReply = &message
|
|
if err := tx.Model(parent).Update("selected_reply_id", message.ID).Error; err != nil {
|
|
return err
|
|
}
|
|
|
|
savedMessages = append(savedMessages, message)
|
|
currentParent = &message
|
|
}
|
|
return nil
|
|
})
|
|
|
|
return savedMessages, err
|
|
}
|
|
|
|
// CloneBranch returns a deep clone of the given message and its replies, returning
|
|
// a new message object. The new message will be attached to the same parent as
|
|
// the messageToClone
|
|
func (s *SQLStore) CloneBranch(messageToClone api.Message) (*api.Message, uint, error) {
|
|
newMessage := messageToClone
|
|
newMessage.ID = 0
|
|
newMessage.Replies = nil
|
|
newMessage.SelectedReplyID = nil
|
|
newMessage.SelectedReply = nil
|
|
|
|
originalReplies, err := s.MessageReplies(messageToClone.ID)
|
|
if err != nil {
|
|
return nil, 0, fmt.Errorf("Could not fetch message %d replies: %v", messageToClone.ID, err)
|
|
}
|
|
|
|
if err := s.db.Create(&newMessage).Error; err != nil {
|
|
return nil, 0, fmt.Errorf("Could not clone message: %s", err)
|
|
}
|
|
|
|
var replyCount uint = 0
|
|
for _, reply := range originalReplies {
|
|
replyCount++
|
|
|
|
newReply := reply
|
|
newReply.ConversationID = messageToClone.ConversationID
|
|
newReply.ParentID = &newMessage.ID
|
|
newReply.Parent = &newMessage
|
|
|
|
res, c, err := s.CloneBranch(newReply)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
newMessage.Replies = append(newMessage.Replies, *res)
|
|
replyCount += c
|
|
|
|
if reply.ID == *messageToClone.SelectedReplyID {
|
|
newMessage.SelectedReplyID = &res.ID
|
|
if err := s.UpdateMessage(&newMessage); err != nil {
|
|
return nil, 0, fmt.Errorf("Could not update parent select reply ID: %v", err)
|
|
}
|
|
}
|
|
}
|
|
return &newMessage, replyCount, nil
|
|
}
|
|
|
|
func fetchMessages(db *gorm.DB) ([]api.Message, error) {
|
|
var messages []api.Message
|
|
if err := db.Preload("Conversation").Find(&messages).Error; err != nil {
|
|
return nil, fmt.Errorf("Could not fetch messages: %v", err)
|
|
}
|
|
|
|
messageMap := make(map[uint]api.Message)
|
|
for i, message := range messages {
|
|
messageMap[messages[i].ID] = message
|
|
}
|
|
|
|
// Create a map to store replies by their parent ID
|
|
repliesMap := make(map[uint][]api.Message)
|
|
for i, message := range messages {
|
|
if messages[i].ParentID != nil {
|
|
repliesMap[*messages[i].ParentID] = append(repliesMap[*messages[i].ParentID], message)
|
|
}
|
|
}
|
|
|
|
// Assign replies, parent, and selected reply to each message
|
|
for i := range messages {
|
|
if replies, exists := repliesMap[messages[i].ID]; exists {
|
|
messages[i].Replies = make([]api.Message, len(replies))
|
|
for j, m := range replies {
|
|
messages[i].Replies[j] = m
|
|
}
|
|
}
|
|
if messages[i].ParentID != nil {
|
|
if parent, exists := messageMap[*messages[i].ParentID]; exists {
|
|
messages[i].Parent = &parent
|
|
}
|
|
}
|
|
if messages[i].SelectedReplyID != nil {
|
|
if selectedReply, exists := messageMap[*messages[i].SelectedReplyID]; exists {
|
|
messages[i].SelectedReply = &selectedReply
|
|
}
|
|
}
|
|
}
|
|
return messages, nil
|
|
}
|
|
|
|
func (s *SQLStore) buildPath(message *api.Message, getNext func(*api.Message) *uint) ([]api.Message, error) {
|
|
var messages []api.Message
|
|
messages, err := fetchMessages(s.db.Where("conversation_id = ?", message.ConversationID))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Create a map to store messages by their ID
|
|
messageMap := make(map[uint]*api.Message)
|
|
for i := range messages {
|
|
messageMap[messages[i].ID] = &messages[i]
|
|
}
|
|
|
|
// Build the path
|
|
var path []api.Message
|
|
nextID := &message.ID
|
|
|
|
for {
|
|
current, exists := messageMap[*nextID]
|
|
if !exists {
|
|
return nil, fmt.Errorf("Message with ID %d not found in conversation", *nextID)
|
|
}
|
|
|
|
path = append(path, *current)
|
|
|
|
nextID = getNext(current)
|
|
if nextID == nil {
|
|
break
|
|
}
|
|
}
|
|
|
|
return path, nil
|
|
}
|
|
|
|
// PathToRoot traverses the provided message's Parent until reaching the tree
|
|
// root and returns a slice of all messages traversed in chronological order
|
|
// (starting with the root and ending with the message provided)
|
|
func (s *SQLStore) PathToRoot(message *api.Message) ([]api.Message, error) {
|
|
if message == nil || message.ID <= 0 {
|
|
return nil, fmt.Errorf("Message is nil or has invalid ID")
|
|
}
|
|
|
|
path, err := s.buildPath(message, func(m *api.Message) *uint {
|
|
return m.ParentID
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
slices.Reverse(path)
|
|
|
|
return path, nil
|
|
}
|
|
|
|
// PathToLeaf traverses the provided message's SelectedReply until reaching a
|
|
// tree leaf and returns a slice of all messages traversed in chronological
|
|
// order (starting with the message provided and ending with the leaf)
|
|
func (s *SQLStore) PathToLeaf(message *api.Message) ([]api.Message, error) {
|
|
if message == nil || message.ID <= 0 {
|
|
return nil, fmt.Errorf("Message is nil or has invalid ID")
|
|
}
|
|
|
|
return s.buildPath(message, func(m *api.Message) *uint {
|
|
return m.SelectedReplyID
|
|
})
|
|
}
|
|
|
|
func (s *SQLStore) LatestConversationMessages() ([]api.Message, error) {
|
|
var latestMessages []api.Message
|
|
|
|
subQuery := s.db.Model(&api.Message{}).
|
|
Select("MAX(created_at) as max_created_at, conversation_id").
|
|
Group("conversation_id")
|
|
|
|
err := s.db.Model(&api.Message{}).
|
|
Joins("JOIN (?) as sub on messages.conversation_id = sub.conversation_id AND messages.created_at = sub.max_created_at", subQuery).
|
|
Group("messages.conversation_id").
|
|
Order("created_at DESC").
|
|
Preload("Conversation.SelectedRoot").
|
|
Find(&latestMessages).Error
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return latestMessages, nil
|
|
}
|