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 MessageRole // one of: 'system', '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) DeleteMessage(message *Message) error { return s.db.Delete(&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 }