Skip to content

Commit

Permalink
robot: switch to kvbrain
Browse files Browse the repository at this point in the history
  • Loading branch information
zephyrtronium committed Mar 27, 2024
1 parent d66156a commit 3fcf203
Show file tree
Hide file tree
Showing 5 changed files with 28 additions and 45 deletions.
4 changes: 2 additions & 2 deletions command/robot.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ package command
import (
"log/slog"

"github.com/zephyrtronium/robot/brain/sqlbrain"
"github.com/zephyrtronium/robot/brain/kvbrain"
"github.com/zephyrtronium/robot/channel"
"github.com/zephyrtronium/robot/privacy"
)
Expand All @@ -12,6 +12,6 @@ import (
type Robot struct {
Log *slog.Logger
Channels map[string]*channel.Channel // TODO(zeph): syncmap[string]channel.Channel
Brain *sqlbrain.Brain
Brain *kvbrain.Brain
Privacy *privacy.List
}
33 changes: 15 additions & 18 deletions config.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ import (
"time"

"github.com/BurntSushi/toml"
"github.com/dgraph-io/badger/v4"
"github.com/dgraph-io/badger/v4/options"
"gitlab.com/zephyrtronium/pick"
"gitlab.com/zephyrtronium/sq"
"gitlab.com/zephyrtronium/tmi"
Expand All @@ -20,7 +22,7 @@ import (
"golang.org/x/time/rate"

"github.com/zephyrtronium/robot/auth"
"github.com/zephyrtronium/robot/brain/sqlbrain"
"github.com/zephyrtronium/robot/brain/kvbrain"
"github.com/zephyrtronium/robot/channel"
"github.com/zephyrtronium/robot/message"
"github.com/zephyrtronium/robot/privacy"
Expand Down Expand Up @@ -60,12 +62,9 @@ func (robo *Robot) SetSecrets(file string) error {

// SetSources opens the brain and privacy list wrappers around the respective
// databases. Use [loadDBs] to open the databases themselves from DSNs.
func (robo *Robot) SetSources(ctx context.Context, brain, priv *sq.DB) error {
func (robo *Robot) SetSources(ctx context.Context, brain *badger.DB, priv *sq.DB) error {
var err error
robo.brain, err = sqlbrain.Open(ctx, brain)
if err != nil {
return fmt.Errorf("couldn't open brain: %w", err)
}
robo.brain = kvbrain.New(brain)
robo.privacy, err = privacy.Open(ctx, priv)
if err != nil {
return fmt.Errorf("couldn't open privacy list: %w", err)
Expand Down Expand Up @@ -154,19 +153,16 @@ func (robo *Robot) SetTwitchChannels(ctx context.Context, global Global, channel
return nil
}

func loadDBs(ctx context.Context, cfg DBCfg) (brain, priv *sq.DB, err error) {
brain, err = sq.Open("sqlite3", cfg.Brain)
func loadDBs(ctx context.Context, cfg DBCfg) (brain *badger.DB, priv *sq.DB, err error) {
opts := badger.DefaultOptions(cfg.Brain)
// TODO(zeph): logger?
opts = opts.WithLogger(nil)
opts = opts.WithCompression(options.None)
opts = opts.WithBloomFalsePositive(0)
brain, err = badger.Open(opts.FromSuperFlag(cfg.Brainflag))
if err != nil {
return nil, nil, fmt.Errorf("couldn't open brain db: %w", err)
}
if err := brain.Ping(ctx); err != nil {
return nil, nil, fmt.Errorf("couldn't connect to brain db: %w", err)
}

if cfg.Privacy == cfg.Brain {
priv = brain
return brain, priv, nil
}

priv, err = sq.Open("sqlite3", cfg.Privacy)
if err != nil {
Expand Down Expand Up @@ -331,8 +327,9 @@ type ClientCfg struct {

// DBCfg is the configuration of databases.
type DBCfg struct {
Brain string `toml:"brain"`
Privacy string `toml:"privacy"`
Brain string `toml:"brain"`
Brainflag string `toml:"brainflag"`
Privacy string `toml:"privacy"`
}

// Rate is a rate limit configuration.
Expand Down
10 changes: 7 additions & 3 deletions example.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,14 @@ name = 'zephyrtronium'
# contact is a description of how to contact the owner.
contact = '/w zephyrtronium'

# db is a table of databases used by the bot. Each one may be the same as or
# different from any other. Each value is a connection string for SQLite3.
# db is a table of databases used by the bot.
[db]
brain = 'file:$ROBOT_SQLITE'
# brain is the directory in which learned knowledge is stored.
brain = '$ROBOT_KNOWLEDGE'
# brainflag configures the brain database as a Badger "superflag" string.
brainflag = ''
# privacy is an SQLite3 connection string for the database where privacy
# information is stored.
privacy = 'file:$ROBOT_SQLITE'

# global includes chat settings that apply to all channels.
Expand Down
20 changes: 1 addition & 19 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ import (
"strings"

"github.com/urfave/cli/v3"
"github.com/zephyrtronium/robot/brain/sqlbrain"
"github.com/zephyrtronium/robot/privacy"
)

Expand All @@ -27,7 +26,6 @@ var app = cli.Command{
Action: cliInit,
Flags: []cli.Flag{
&flagConfig,
&flagOrder,
},
},
{
Expand Down Expand Up @@ -64,7 +62,6 @@ func main() {

func cliInit(ctx context.Context, cmd *cli.Command) error {
slog.SetDefault(loggerFromFlags(cmd))
order := int(cmd.Int("order"))
r, err := os.Open(cmd.String("config"))
if err != nil {
return fmt.Errorf("couldn't open config file: %w", err)
Expand All @@ -73,13 +70,10 @@ func cliInit(ctx context.Context, cmd *cli.Command) error {
if err != nil {
return fmt.Errorf("couldn't load config: %w", err)
}
brain, priv, err := loadDBs(ctx, cfg.DB)
_, priv, err := loadDBs(ctx, cfg.DB)
if err != nil {
return err
}
if err := sqlbrain.Create(ctx, brain, order); err != nil {
return fmt.Errorf("couldn't initialize brain: %w", err)
}
if err := privacy.Init(ctx, priv); err != nil {
return fmt.Errorf("couldn't initialize privacy list: %w", err)
}
Expand Down Expand Up @@ -162,18 +156,6 @@ var (
}
},
}

flagOrder = cli.IntFlag{
Name: "order",
Required: true,
Usage: "Prefix length for Markov chains",
Action: func(ctx context.Context, cmd *cli.Command, i int64) error {
if i <= 0 {
return errors.New("order must be positive")
}
return nil
},
}
)

func loggerFromFlags(cmd *cli.Command) *slog.Logger {
Expand Down
6 changes: 3 additions & 3 deletions robot.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,15 +18,15 @@ import (
"golang.org/x/time/rate"

"github.com/zephyrtronium/robot/auth"
"github.com/zephyrtronium/robot/brain/sqlbrain"
"github.com/zephyrtronium/robot/brain/kvbrain"
"github.com/zephyrtronium/robot/channel"
"github.com/zephyrtronium/robot/privacy"
)

// Robot is the overall configuration for the bot.
type Robot struct {
// brain is the brain.
brain *sqlbrain.Brain
brain *kvbrain.Brain
// privacy is the privacy.
privacy *privacy.List
// channels are the channels.
Expand Down Expand Up @@ -164,7 +164,7 @@ func validateTwitch(ctx context.Context, tok *oauth2.Token) error {
if err := json.Unmarshal(body, &s); err != nil {
return fmt.Errorf("couldn't unmarshal token validation response: %w", err)
}
slog.InfoContext(ctx, "token validation", slog.String("token", tok.AccessToken), slog.Any("result", s))
slog.InfoContext(ctx, "token validation", slog.Any("result", s))
if resp.StatusCode == http.StatusUnauthorized {
// Token expired or otherwise invalid. We need a refresh.
return fmt.Errorf("token validation failed: %s (%w)", s.Message, errNeedRefresh)
Expand Down

0 comments on commit 3fcf203

Please sign in to comment.