Large refactor - it compiles!
This refactor splits out all conversation concerns into a new `conversation` package. There is now a split between `conversation` and `api`s representation of `Message`, the latter storing the minimum information required for interaction with LLM providers. There is necessary conversation between the two when making LLM calls.
This commit is contained in:
@@ -12,11 +12,12 @@ import (
|
||||
|
||||
"git.mlow.ca/mlow/lmcli/pkg/agents"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/api"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/api/provider"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/api/provider/anthropic"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/api/provider/google"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/api/provider/ollama"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/api/provider/openai"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/provider"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/provider/anthropic"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/provider/google"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/provider/ollama"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/provider/openai"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/conversation"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/util"
|
||||
"git.mlow.ca/mlow/lmcli/pkg/util/tty"
|
||||
"gorm.io/driver/sqlite"
|
||||
@@ -33,7 +34,7 @@ type Agent struct {
|
||||
type Context struct {
|
||||
// high level app configuration, may be mutated at runtime
|
||||
Config Config
|
||||
Store ConversationStore
|
||||
Conversations conversation.Repo
|
||||
Chroma *tty.ChromaHighlighter
|
||||
}
|
||||
|
||||
@@ -44,7 +45,7 @@ func NewContext() (*Context, error) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
store, err := getConversationStore()
|
||||
store, err := getConversationService()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -69,17 +70,16 @@ func createOrOpenAppend(path string) (*os.File, error) {
|
||||
return file, nil
|
||||
}
|
||||
|
||||
func getConversationStore() (ConversationStore, error) {
|
||||
func getConversationService() (conversation.Repo, error) {
|
||||
databaseFile := filepath.Join(dataDir(), "conversations.db")
|
||||
|
||||
gormLogFile, err := createOrOpenAppend(filepath.Join(dataDir(), "database.log"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Could not open database log file: %v", err)
|
||||
}
|
||||
db, err := gorm.Open(sqlite.Open(databaseFile), &gorm.Config{
|
||||
Logger: logger.New(log.New(gormLogFile, "", log.LstdFlags), logger.Config{
|
||||
Logger: logger.New(log.New(gormLogFile, "\n", log.LstdFlags), logger.Config{
|
||||
SlowThreshold: 200 * time.Millisecond,
|
||||
LogLevel: logger.Warn,
|
||||
LogLevel: logger.Info,
|
||||
IgnoreRecordNotFoundError: false,
|
||||
Colorful: true,
|
||||
}),
|
||||
@@ -87,11 +87,11 @@ func getConversationStore() (ConversationStore, error) {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Error establishing connection to store: %v", err)
|
||||
}
|
||||
store, err := NewSQLStore(db)
|
||||
repo, err := conversation.NewRepo(db)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return store, nil
|
||||
return repo, nil
|
||||
}
|
||||
|
||||
func (c *Context) GetModels() (models []string) {
|
||||
|
||||
@@ -1,433 +0,0 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user