Add initial store.go for conversation/message persistence
This commit is contained in:
parent
fd667a1f47
commit
8fe2a2cf53
6
go.mod
6
go.mod
@ -5,9 +5,15 @@ go 1.19
|
|||||||
require (
|
require (
|
||||||
github.com/sashabaranov/go-openai v1.16.0
|
github.com/sashabaranov/go-openai v1.16.0
|
||||||
github.com/spf13/cobra v1.7.0
|
github.com/spf13/cobra v1.7.0
|
||||||
|
github.com/sqids/sqids-go v0.4.1
|
||||||
|
gorm.io/driver/sqlite v1.5.4
|
||||||
|
gorm.io/gorm v1.25.5
|
||||||
)
|
)
|
||||||
|
|
||||||
require (
|
require (
|
||||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||||
|
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||||
|
github.com/jinzhu/now v1.1.5 // indirect
|
||||||
|
github.com/mattn/go-sqlite3 v1.14.17 // indirect
|
||||||
github.com/spf13/pflag v1.0.5 // indirect
|
github.com/spf13/pflag v1.0.5 // indirect
|
||||||
)
|
)
|
||||||
|
12
go.sum
12
go.sum
@ -1,6 +1,12 @@
|
|||||||
github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
|
github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
|
||||||
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
|
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
|
||||||
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
|
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
|
||||||
|
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
|
||||||
|
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
|
||||||
|
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
|
||||||
|
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
|
||||||
|
github.com/mattn/go-sqlite3 v1.14.17 h1:mCRHCLDUBXgpKAqIKsaAaAsrAlbkeomtRFKXh2L6YIM=
|
||||||
|
github.com/mattn/go-sqlite3 v1.14.17/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
|
||||||
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
|
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
|
||||||
github.com/sashabaranov/go-openai v1.16.0 h1:34W6WV84ey6OpW0p2UewZkdMu82AxGC+BzpU6iiauRw=
|
github.com/sashabaranov/go-openai v1.16.0 h1:34W6WV84ey6OpW0p2UewZkdMu82AxGC+BzpU6iiauRw=
|
||||||
github.com/sashabaranov/go-openai v1.16.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
|
github.com/sashabaranov/go-openai v1.16.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
|
||||||
@ -8,5 +14,11 @@ github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
|
|||||||
github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0=
|
github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0=
|
||||||
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
|
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
|
||||||
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
|
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
|
||||||
|
github.com/sqids/sqids-go v0.4.1 h1:eQKYzmAZbLlRwHeHYPF35QhgxwZHLnlmVj9AkIj/rrw=
|
||||||
|
github.com/sqids/sqids-go v0.4.1/go.mod h1:EMwHuPQgSNFS0A49jESTfIQS+066XQTVhukrzEPScl8=
|
||||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||||
|
gorm.io/driver/sqlite v1.5.4 h1:IqXwXi8M/ZlPzH/947tn5uik3aYQslP9BVveoax0nV0=
|
||||||
|
gorm.io/driver/sqlite v1.5.4/go.mod h1:qxAuCol+2r6PannQDpOP1FP6ag3mKi4esLnB/jHed+4=
|
||||||
|
gorm.io/gorm v1.25.5 h1:zR9lOiiYf09VNh5Q1gphfyia1JpiClIWG9hQaxB/mls=
|
||||||
|
gorm.io/gorm v1.25.5/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8=
|
||||||
|
20
main.go
20
main.go
@ -7,21 +7,13 @@ import (
|
|||||||
"github.com/spf13/cobra"
|
"github.com/spf13/cobra"
|
||||||
)
|
)
|
||||||
|
|
||||||
type Message struct {
|
var store, storeError = InitializeStore()
|
||||||
MessageID string
|
|
||||||
ConversationID string
|
|
||||||
Conversation Conversation
|
|
||||||
OriginalContent string
|
|
||||||
Role string // 'user' or 'assistant'
|
|
||||||
}
|
|
||||||
|
|
||||||
type Conversation struct {
|
func checkStore() {
|
||||||
ID string
|
if storeError != nil {
|
||||||
Title string
|
fmt.Fprintf(os.Stderr, "Error establishing connection to store: %v\n", storeError)
|
||||||
|
os.Exit(1)
|
||||||
}
|
}
|
||||||
|
|
||||||
type Context struct {
|
|
||||||
CurrentConversation string
|
|
||||||
}
|
}
|
||||||
|
|
||||||
var rootCmd = &cobra.Command{
|
var rootCmd = &cobra.Command{
|
||||||
@ -29,6 +21,7 @@ var rootCmd = &cobra.Command{
|
|||||||
Short: "Interact with Large Language Models",
|
Short: "Interact with Large Language Models",
|
||||||
Long: `lm is a CLI tool to interact with OpenAI's GPT 3.5 and GPT 4.`,
|
Long: `lm is a CLI tool to interact with OpenAI's GPT 3.5 and GPT 4.`,
|
||||||
Run: func(cmd *cobra.Command, args []string) {
|
Run: func(cmd *cobra.Command, args []string) {
|
||||||
|
checkStore()
|
||||||
// execute `lm ls` by default
|
// execute `lm ls` by default
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
@ -100,7 +93,6 @@ var newCmd = &cobra.Command{
|
|||||||
|
|
||||||
fmt.Printf("> %s\n", messageContents)
|
fmt.Printf("> %s\n", messageContents)
|
||||||
|
|
||||||
// Initialize the messages array for this conversation.
|
|
||||||
messages := []Message{
|
messages := []Message{
|
||||||
{
|
{
|
||||||
OriginalContent: messageContents,
|
OriginalContent: messageContents,
|
||||||
|
87
store.go
Normal file
87
store.go
Normal file
@ -0,0 +1,87 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"gorm.io/gorm"
|
||||||
|
"gorm.io/driver/sqlite"
|
||||||
|
sqids "github.com/sqids/sqids-go"
|
||||||
|
)
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
const (
|
||||||
|
DATABASE_FILE string = "./data.db"
|
||||||
|
)
|
||||||
|
|
||||||
|
func InitializeStore() (*Store, error) {
|
||||||
|
db, err := gorm.Open(sqlite.Open(DATABASE_FILE), &gorm.Config{})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
models := []any{
|
||||||
|
&Conversation{},
|
||||||
|
&Message{},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, x := range(models) {
|
||||||
|
err := db.AutoMigrate(x)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
_sqids, _ := sqids.New(sqids.Options{
|
||||||
|
MinLength: 4,
|
||||||
|
})
|
||||||
|
|
||||||
|
return &Store{db: db, sqids: _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) 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
|
||||||
|
}
|
Loading…
Reference in New Issue
Block a user