package cli import ( "database/sql" "os" "path/filepath" 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 // 'user' or 'assistant' } type Conversation struct { ID uint `gorm:"primaryKey"` ShortName sql.NullString Title string } func getDataDir() 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 InitializeStore() *Store { databaseFile := filepath.Join(getDataDir(), "conversations.db") db, err := gorm.Open(sqlite.Open(databaseFile), &gorm.Config{}) if err != nil { Fatal("Error establishing connection to store: %v\n", err) return nil } models := []any{ &Conversation{}, &Message{}, } for _, x := range models { err := db.AutoMigrate(x) if err != nil { Fatal("Could not perform database migrations: %v\n", err) return nil } } _sqids, _ := sqids.New(sqids.Options{MinLength: 4}) return &Store{db, _sqids} } 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) SaveMessage(message *Message) error { return s.db.Create(message).Error } func (s *Store) GetConversations() ([]Conversation, error) { var conversations []Conversation err := s.db.Find(&conversations).Error return conversations, err } func (s *Store) GetMessages(conversation *Conversation) ([]Message, error) { var messages []Message err := s.db.Where("conversation_id = ?", conversation.ID).Find(&messages).Error return messages, err }