lmcli/pkg/cli/store.go

package cli

import (
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	sqids "github.com/sqids/sqids-go"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

type Store struct {
	db    *gorm.DB
	sqids *sqids.Sqids
}

type Message struct {
	ID              uint `gorm:"primaryKey"`
	ConversationID  uint `gorm:"foreignKey:ConversationID"`
	Conversation    Conversation
	OriginalContent string
	Role            string // one of: 'user', 'assistant', 'tool'
	CreatedAt       time.Time
	ToolCallID		sql.NullString
	ToolCalls		sql.NullString // a json-encoded array of tool calls from the model
}

type Conversation struct {
	ID        uint `gorm:"primaryKey"`
	ShortName sql.NullString
	Title     string
}

func DataDir() string {
	var dataDir string

	xdgDataHome := os.Getenv("XDG_DATA_HOME")
	if xdgDataHome != "" {
		dataDir = filepath.Join(xdgDataHome, "lmcli")
	} else {
		userHomeDir, _ := os.UserHomeDir()
		dataDir = filepath.Join(userHomeDir, ".local/share/lmcli")
	}

	os.MkdirAll(dataDir, 0755)
	return dataDir
}

func NewStore() (*Store, error) {
	databaseFile := filepath.Join(DataDir(), "conversations.db")
	db, err := gorm.Open(sqlite.Open(databaseFile), &gorm.Config{})
	if err != nil {
		return nil, fmt.Errorf("Error establishing connection to store: %v", err)
	}

	models := []any{
		&Conversation{},
		&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 &Store{db, _sqids}, nil
}

func (s *Store) SaveConversation(conversation *Conversation) error {
	err := s.db.Save(&conversation).Error
	if err != nil {
		return err
	}

	if !conversation.ShortName.Valid {
		shortName, _ := s.sqids.Encode([]uint64{uint64(conversation.ID)})
		conversation.ShortName = sql.NullString{String: shortName, Valid: true}
		err = s.db.Save(&conversation).Error
	}

	return err
}

func (s *Store) DeleteConversation(conversation *Conversation) error {
	s.db.Where("conversation_id = ?", conversation.ID).Delete(&Message{})
	return s.db.Delete(&conversation).Error
}

func (s *Store) SaveMessage(message *Message) error {
	return s.db.Create(message).Error
}

func (s *Store) Conversations() ([]Conversation, error) {
	var conversations []Conversation
	err := s.db.Find(&conversations).Error
	return conversations, err
}

func (s *Store) ConversationShortNameCompletions(shortName string) []string {
	var completions []string
	conversations, _ := s.Conversations() // ignore error for completions
	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 *Store) ConversationByShortName(shortName string) (*Conversation, error) {
	var conversation Conversation
	err := s.db.Where("short_name = ?", shortName).Find(&conversation).Error
	return &conversation, err
}

func (s *Store) Messages(conversation *Conversation) ([]Message, error) {
	var messages []Message
	err := s.db.Where("conversation_id = ?", conversation.ID).Find(&messages).Error
	return messages, err
}

func (s *Store) LastMessage(conversation *Conversation) (*Message, error) {
	var message Message
	err := s.db.Where("conversation_id = ?", conversation.ID).Last(&message).Error
	return &message, err
}