decouple first step

gpt/gpt.go

@@ -0,0 +1 @@
+package gpt

main.go

@@ -1,59 +1,27 @@
 package main
 
 import (
-	"context"
-	"errors"
-	"fmt"
-	"github.com/chzyer/readline"
-	_ "github.com/mattn/go-sqlite3"
-	"github.com/rs/zerolog"
-	"github.com/rs/zerolog/log"
-	"github.com/sashabaranov/go-openai"
-	"golang.org/x/exp/slog"
-	"maunium.net/go/mautrix"
-	"maunium.net/go/mautrix/crypto/cryptohelper"
-	"maunium.net/go/mautrix/event"
-	"maunium.net/go/mautrix/format"
-	"maunium.net/go/mautrix/id"
 	"os"
 	"os/signal"
-	"sync"
+
-	"time"
+	"ewintr.nl/matrix-bots/matrix"
+	_ "github.com/mattn/go-sqlite3"
+	"github.com/sashabaranov/go-openai"
+	"golang.org/x/exp/slog"
 )
 
 func main() {
 	logger := slog.New(slog.NewTextHandler(os.Stderr, nil))
 
-	MatrixHomeserver, ok := os.LookupEnv("MATRIX_HOMESERVER")
+	matrixClient := matrix.New(matrix.Config{
-	if !ok {
+		Homeserver:    getParam("MATRIX_HOMESERVER", "http://localhost"),
-		logger.Error("MATRIX_HOME_SERVER is not set!")
+		UserID:        getParam("MATRIX_USER_ID", "@bot:localhost"),
-		os.Exit(1)
+		UserPassword:  getParam("MATRIX_PASSWORD", "secret"),
-	}
+		UserAccessKey: getParam("MATRIX_ACCESS_KEY", "secret"),
-	MatrixUserIDStr, ok := os.LookupEnv("MATRIX_USER_ID")
+		DBPath:        getParam("BOT_DB_PATH", "bot.db"),
-	if !ok {
+		Pickle:        getParam("BOT_PICKLE", "scrambled"),
-		logger.Error("MATRIX_USER_ID is not set")
+	})
-		os.Exit(1)
+
-	}
-	MatrixUserID := id.UserID(MatrixUserIDStr)
-	MatrixAccessKey, ok := os.LookupEnv("MATRIX_ACCESS_KEY")
-	if !ok {
-		logger.Error("MATRIX_ACCESS_KEY is not set")
-		os.Exit(1)
-	}
-	MatrixPassword, ok := os.LookupEnv("MATRIX_PASSWORD")
-	if !ok {
-		logger.Error("MATRIX_PASSWORD is not set")
-		os.Exit(1)
-	}
-	BotDatabasePath, ok := os.LookupEnv("BOT_DATABASE_PATH")
-	if !ok {
-		logger.Error("BOT_DATABASE_PATH is not set")
-		os.Exit(1)
-	}
-	BotPickle, ok := os.LookupEnv("BOT_PICKLE")
-	if !ok {
-		logger.Error("BOT_PICKLE is not set")
-	}
 	OpenaiAPIKey, ok := os.LookupEnv("OPENAI_API_KEY")
 	if !ok {
 		logger.Error("OPENAI_API_KEY is not set")
@@ -63,139 +31,25 @@ func main() {
 	// Create new OpenAI client
 	openaiClient := openai.NewClient(OpenaiAPIKey)
 
-	// Create new Matrix client
+	if err := matrixClient.Init(); err != nil {
-	client, err := mautrix.NewClient(MatrixHomeserver, MatrixUserID, MatrixAccessKey)
+		logger.Error(err.Error())
-	if err != nil {
+		os.Exit(1)
-		panic(err)
 	}
+	go matrixClient.Run()
 
-	var oei mautrix.OldEventIgnorer
+	matrixClient.AddEventHandler(matrixClient.InviteHandler())
-	oei.Register(client.Syncer.(mautrix.ExtensibleSyncer))
+	matrixClient.AddEventHandler(matrixClient.RespondHandler(openaiClient))
-
-	// logger for matrix client
-	rl, err := readline.New("[no room]> ")
-	if err != nil {
-		panic(err)
-	}
-	defer rl.Close()
-	matrixLogger := zerolog.New(zerolog.NewConsoleWriter(func(w *zerolog.ConsoleWriter) {
-		w.Out = rl.Stdout()
-		w.TimeFormat = time.Stamp
-	})).With().Timestamp().Logger()
-	client.Log = matrixLogger
-
-	var lastRoomID id.RoomID
-
-	// Set up sync events
-	syncer := client.Syncer.(*mautrix.DefaultSyncer)
-
-	syncer.OnEventType(event.EventMessage, func(source mautrix.EventSource, evt *event.Event) {
-		lastRoomID = evt.RoomID
-		rl.SetPrompt(fmt.Sprintf("%s> ", lastRoomID))
-		matrixLogger.Info().
-			Str("sender", evt.Sender.String()).
-			Str("type", evt.Type.String()).
-			Str("id", evt.ID.String()).
-			Str("body", evt.Content.AsMessage().Body).
-			Msg("Received message")
-
-		if evt.Sender != MatrixUserID {
-			msgBody := evt.Content.AsMessage().Body
-
-			// Generate a message with OpenAI API
-			openAiResp, err := openaiClient.CreateChatCompletion(
-				context.Background(),
-				openai.ChatCompletionRequest{
-					Model: openai.GPT4,
-					Messages: []openai.ChatCompletionMessage{
-						{
-							Role:    openai.ChatMessageRoleSystem,
-							Content: "You are a chatbot that helps people by responding to their questions with short messages.",
-						},
-
-						{
-							Role:    openai.ChatMessageRoleUser,
-							Content: msgBody,
-						},
-					},
-				})
-
-			if err != nil {
-				fmt.Println("OpenAI API returned with ", err)
-				return
-			}
-
-			// Send the OpenAI response back to the chat
-			responseMarkdown := openAiResp.Choices[len(openAiResp.Choices)-1].Message.Content
-			responseMessage := format.RenderMarkdown(responseMarkdown, true, false)
-			client.SendMessageEvent(lastRoomID, event.EventMessage, &responseMessage)
-		}
-
-	})
-	syncer.OnEventType(event.StateMember, func(source mautrix.EventSource, evt *event.Event) {
-		if evt.GetStateKey() == client.UserID.String() && evt.Content.AsMember().Membership == event.MembershipInvite {
-			_, err := client.JoinRoomByID(evt.RoomID)
-			if err == nil {
-				lastRoomID = evt.RoomID
-				rl.SetPrompt(fmt.Sprintf("%s> ", lastRoomID))
-				log.Info().
-					Str("room_id", evt.RoomID.String()).
-					Str("inviter", evt.Sender.String()).
-					Msg("Joined room after invite")
-			} else {
-				log.Error().Err(err).
-					Str("room_id", evt.RoomID.String()).
-					Str("inviter", evt.Sender.String()).
-					Msg("Failed to join room after invite")
-			}
-		}
-	})
-
-	cryptoHelper, err := cryptohelper.NewCryptoHelper(client, []byte(BotPickle), BotDatabasePath)
-	if err != nil {
-		panic(err)
-	}
-
-	// You can also store the user/device IDs and access token and put them in the client beforehand instead of using LoginAs.
-	//client.UserID = "..."
-	//client.DeviceID = "..."
-	//client.AccessToken = "..."
-	// You don't need to set a device ID in LoginAs because the crypto helper will set it for you if necessary.
-	cryptoHelper.LoginAs = &mautrix.ReqLogin{
-		Type:       mautrix.AuthTypePassword,
-		Identifier: mautrix.UserIdentifier{Type: mautrix.IdentifierTypeUser, User: MatrixUserID.String()},
-		Password:   MatrixPassword,
-	}
-	// If you want to use multiple clients with the same DB, you should set a distinct database account ID for each one.
-	//cryptoHelper.DBAccountID = ""
-	err = cryptoHelper.Init()
-	if err != nil {
-		panic(err)
-	}
-	// Set the client crypto helper in order to automatically encrypt outgoing messages
-	client.Crypto = cryptoHelper
-
-	matrixLogger.Info().Msg("Now running")
-	syncCtx, cancelSync := context.WithCancel(context.Background())
-	var syncStopWait sync.WaitGroup
-	syncStopWait.Add(1)
-
-	go func() {
-		err = client.SyncWithContext(syncCtx)
-		defer syncStopWait.Done()
-		if err != nil && !errors.Is(err, context.Canceled) {
-			panic(err)
-		}
-	}()
 
 	done := make(chan os.Signal)
 	signal.Notify(done, os.Interrupt)
 	<-done
 
-	cancelSync()
+}
-	syncStopWait.Wait()
+
-	err = cryptoHelper.Close()
+func getParam(name, def string) string {
-	if err != nil {
+	val, ok := os.LookupEnv(name)
-		log.Error().Err(err).Msg("Error closing database")
+	if !ok {
-	}
+		return def
+	}
+	return val
 }

matrix/matrix.go

@@ -0,0 +1,164 @@
+package matrix
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/chzyer/readline"
+	"github.com/rs/zerolog"
+	"github.com/sashabaranov/go-openai"
+	"maunium.net/go/mautrix"
+	"maunium.net/go/mautrix/crypto/cryptohelper"
+	"maunium.net/go/mautrix/event"
+	"maunium.net/go/mautrix/format"
+	"maunium.net/go/mautrix/id"
+)
+
+type Config struct {
+	Homeserver    string
+	UserID        string
+	UserAccessKey string
+	UserPassword  string
+	DBPath        string
+	Pickle        string
+}
+
+type Matrix struct {
+	config       Config
+	readline     *readline.Instance
+	client       *mautrix.Client
+	cryptoHelper *cryptohelper.CryptoHelper
+}
+
+func New(cfg Config) *Matrix {
+	return &Matrix{
+		config: cfg,
+	}
+}
+
+func (m *Matrix) Init() error {
+	rl, err := readline.New("[no room]> ")
+	if err != nil {
+		return err
+	}
+	m.readline = rl
+	defer m.readline.Close()
+
+	client, err := mautrix.NewClient(m.config.Homeserver, id.UserID(m.config.UserID), m.config.UserAccessKey)
+	if err != nil {
+		return err
+	}
+	var oei mautrix.OldEventIgnorer
+	oei.Register(client.Syncer.(mautrix.ExtensibleSyncer))
+	m.client = client
+
+	m.client.Log = zerolog.New(zerolog.NewConsoleWriter(func(w *zerolog.ConsoleWriter) {
+		w.Out = rl.Stdout()
+		w.TimeFormat = time.Stamp
+	})).With().Timestamp().Logger()
+
+	m.cryptoHelper, err = cryptohelper.NewCryptoHelper(client, []byte(m.config.Pickle), m.config.DBPath)
+	if err != nil {
+		return err
+	}
+	m.cryptoHelper.LoginAs = &mautrix.ReqLogin{
+		Type:       mautrix.AuthTypePassword,
+		Identifier: mautrix.UserIdentifier{Type: mautrix.IdentifierTypeUser, User: m.config.UserID},
+		Password:   m.config.UserPassword,
+	}
+	if err := m.cryptoHelper.Init(); err != nil {
+		return err
+	}
+	m.client.Crypto = m.cryptoHelper
+
+	return nil
+}
+
+func (m *Matrix) Run() error {
+	if err := m.client.Sync(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (m *Matrix) Close() error {
+	if err := m.client.Sync(); err != nil {
+		return err
+	}
+	if err := m.cryptoHelper.Close(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (m *Matrix) AddEventHandler(eventType event.Type, handler mautrix.EventHandler) {
+	syncer := m.client.Syncer.(*mautrix.DefaultSyncer)
+	syncer.OnEventType(eventType, handler)
+}
+
+func (m *Matrix) InviteHandler() (event.Type, mautrix.EventHandler) {
+	return event.StateMember, func(source mautrix.EventSource, evt *event.Event) {
+		if evt.GetStateKey() == m.client.UserID.String() && evt.Content.AsMember().Membership == event.MembershipInvite {
+			_, err := m.client.JoinRoomByID(evt.RoomID)
+			if err == nil {
+				m.readline.SetPrompt(fmt.Sprintf("%s> ", evt.RoomID))
+				m.client.Log.Info().
+					Str("room_id", evt.RoomID.String()).
+					Str("inviter", evt.Sender.String()).
+					Msg("Joined room after invite")
+			} else {
+				m.client.Log.Error().Err(err).
+					Str("room_id", evt.RoomID.String()).
+					Str("inviter", evt.Sender.String()).
+					Msg("Failed to join room after invite")
+			}
+		}
+	}
+}
+
+func (m *Matrix) RespondHandler(gpt *openai.Client) (event.Type, mautrix.EventHandler) {
+	return event.EventMessage, func(source mautrix.EventSource, evt *event.Event) {
+		m.readline.SetPrompt(fmt.Sprintf("%s> ", evt.RoomID))
+		m.client.Log.Info().
+			Str("sender", evt.Sender.String()).
+			Str("type", evt.Type.String()).
+			Str("id", evt.ID.String()).
+			Str("body", evt.Content.AsMessage().Body).
+			Msg("Received message")
+
+		if evt.Sender != id.UserID(m.config.UserID) {
+			msgBody := evt.Content.AsMessage().Body
+
+			// Generate a message with OpenAI API
+			openAiResp, err := gpt.CreateChatCompletion(
+				context.Background(),
+				openai.ChatCompletionRequest{
+					Model: openai.GPT4,
+					Messages: []openai.ChatCompletionMessage{
+						{
+							Role:    openai.ChatMessageRoleSystem,
+							Content: "You are a chatbot that helps people by responding to their questions with short messages.",
+						},
+
+						{
+							Role:    openai.ChatMessageRoleUser,
+							Content: msgBody,
+						},
+					},
+				})
+
+			if err != nil {
+				fmt.Println("OpenAI API returned with ", err)
+				return
+			}
+
+			// Send the OpenAI response back to the chat
+			responseMarkdown := openAiResp.Choices[len(openAiResp.Choices)-1].Message.Content
+			responseMessage := format.RenderMarkdown(responseMarkdown, true, false)
+			m.client.SendMessageEvent(evt.RoomID, event.EventMessage, &responseMessage)
+		}
+	}
+}