132 lines
3.3 KiB
Go
132 lines
3.3 KiB
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 // 'user' or 'assistant'
|
|
CreatedAt time.Time
|
|
}
|
|
|
|
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
|
|
}
|