diff --git a/gpt/gpt.go b/gpt/gpt.go new file mode 100644 index 0000000..56f2894 --- /dev/null +++ b/gpt/gpt.go @@ -0,0 +1 @@ +package gpt diff --git a/main.go b/main.go index 86e2699..299b8aa 100644 --- a/main.go +++ b/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") - if !ok { - logger.Error("MATRIX_HOME_SERVER is not set!") - os.Exit(1) - } - MatrixUserIDStr, ok := os.LookupEnv("MATRIX_USER_ID") - if !ok { - 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") - } + matrixClient := matrix.New(matrix.Config{ + Homeserver: getParam("MATRIX_HOMESERVER", "http://localhost"), + UserID: getParam("MATRIX_USER_ID", "@bot:localhost"), + UserPassword: getParam("MATRIX_PASSWORD", "secret"), + UserAccessKey: getParam("MATRIX_ACCESS_KEY", "secret"), + DBPath: getParam("BOT_DB_PATH", "bot.db"), + Pickle: getParam("BOT_PICKLE", "scrambled"), + }) + 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 - client, err := mautrix.NewClient(MatrixHomeserver, MatrixUserID, MatrixAccessKey) - if err != nil { - panic(err) + if err := matrixClient.Init(); err != nil { + logger.Error(err.Error()) + os.Exit(1) } + go matrixClient.Run() - var oei mautrix.OldEventIgnorer - oei.Register(client.Syncer.(mautrix.ExtensibleSyncer)) - - // 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) - } - }() + matrixClient.AddEventHandler(matrixClient.InviteHandler()) + matrixClient.AddEventHandler(matrixClient.RespondHandler(openaiClient)) done := make(chan os.Signal) signal.Notify(done, os.Interrupt) <-done - cancelSync() - syncStopWait.Wait() - err = cryptoHelper.Close() - if err != nil { - log.Error().Err(err).Msg("Error closing database") - } +} + +func getParam(name, def string) string { + val, ok := os.LookupEnv(name) + if !ok { + return def + } + return val } diff --git a/matrix/matrix.go b/matrix/matrix.go new file mode 100644 index 0000000..feae2b1 --- /dev/null +++ b/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) + } + } +}