diff --git a/Dockerfile b/Dockerfile index 1719cc0..0daffdd 100644 --- a/Dockerfile +++ b/Dockerfile @@ -4,10 +4,10 @@ RUN apt update && apt install -y libolm3 libolm-dev WORKDIR /src COPY . ./ RUN go mod download -RUN go build -o /matrix-bots ./main.go +RUN go build -o /matrix-gptzoo ./main.go FROM debian:bullseye RUN apt update && apt install -y libolm3 ca-certificates openssl -COPY --from=build /matrix-bots /matrix-bots -CMD /matrix-bots \ No newline at end of file +COPY --from=build /matrix-gptzoo /matrix-gptzoo +CMD /matrix-gptzoo \ No newline at end of file diff --git a/bot/conversation.go b/bot/conversation.go new file mode 100644 index 0000000..57da5cf --- /dev/null +++ b/bot/conversation.go @@ -0,0 +1,69 @@ +package bot + +import "github.com/sashabaranov/go-openai" + +const systemPrompt = "You are a chatbot that helps people by responding to their questions with short messages." + +type Message struct { + EventID string + Role string + Content string + ReplyToID string +} + +type Conversation struct { + Messages []Message +} + +func NewConversation(question string) *Conversation { + return &Conversation{ + Messages: []Message{ + { + Role: openai.ChatMessageRoleSystem, + Content: systemPrompt, + }, + { + Role: openai.ChatMessageRoleUser, + Content: question, + }, + }, + } +} + +func (c *Conversation) Contains(EventID string) bool { + for _, m := range c.Messages { + if m.EventID == EventID { + return true + } + } + + return false +} + +func (c *Conversation) Add(msg Message) { + c.Messages = append(c.Messages, msg) +} + +type Conversations []*Conversation + +func (cs Conversations) Contains(EventID string) bool { + for _, c := range cs { + if c.Contains(EventID) { + return true + } + } + + return false +} + +func (cs Conversations) Add(msg Message) { + for _, c := range cs { + if c.Contains(msg.EventID) { + c.Add(msg) + return + } + } + + c := NewConversation(msg.Content) + cs = append(cs, c) +} diff --git a/bot/gpt.go b/bot/gpt.go new file mode 100644 index 0000000..fc959b4 --- /dev/null +++ b/bot/gpt.go @@ -0,0 +1,39 @@ +package bot + +import ( + "context" + + "github.com/sashabaranov/go-openai" +) + +type GPT struct { + client *openai.Client +} + +func NewGPT(apiKey string) *GPT { + return &GPT{ + client: openai.NewClient(apiKey), + } +} + +func (g GPT) Complete(conv *Conversation) (string, error) { + ctx := context.Background() + msg := []openai.ChatCompletionMessage{} + for _, m := range conv.Messages { + msg = append(msg, openai.ChatCompletionMessage{ + Role: m.Role, + Content: m.Content, + }) + } + req := openai.ChatCompletionRequest{ + Model: openai.GPT4, + Messages: msg, + } + + resp, err := g.client.CreateChatCompletion(ctx, req) + if err != nil { + return "", err + } + + return resp.Choices[len(resp.Choices)-1].Message.Content, nil +} diff --git a/matrix/matrix.go b/bot/matrix.go similarity index 68% rename from matrix/matrix.go rename to bot/matrix.go index feae2b1..5b67ed8 100644 --- a/matrix/matrix.go +++ b/bot/matrix.go @@ -1,13 +1,11 @@ -package matrix +package bot 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" @@ -22,13 +20,16 @@ type Config struct { UserPassword string DBPath string Pickle string + OpenAIKey string } type Matrix struct { - config Config - readline *readline.Instance - client *mautrix.Client - cryptoHelper *cryptohelper.CryptoHelper + config Config + readline *readline.Instance + client *mautrix.Client + cryptoHelper *cryptohelper.CryptoHelper + conversations Conversations + gptClient *GPT } func New(cfg Config) *Matrix { @@ -38,13 +39,6 @@ func New(cfg Config) *Matrix { } 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 @@ -54,7 +48,6 @@ func (m *Matrix) Init() error { 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() @@ -72,6 +65,10 @@ func (m *Matrix) Init() error { } m.client.Crypto = m.cryptoHelper + m.gptClient = NewGPT(m.config.OpenAIKey) + + m.conversations = make(Conversations, 0) + return nil } @@ -104,7 +101,6 @@ func (m *Matrix) InviteHandler() (event.Type, mautrix.EventHandler) { 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()). @@ -119,36 +115,15 @@ func (m *Matrix) InviteHandler() (event.Type, mautrix.EventHandler) { } } -func (m *Matrix) RespondHandler(gpt *openai.Client) (event.Type, mautrix.EventHandler) { +func (m *Matrix) RespondHandler() (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, - }, - }, - }) + resp, err := m.gptClient.Complete(NewConversation(msgBody)) if err != nil { fmt.Println("OpenAI API returned with ", err) @@ -156,9 +131,8 @@ func (m *Matrix) RespondHandler(gpt *openai.Client) (event.Type, mautrix.EventHa } // 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) + formattedResp := format.RenderMarkdown(resp, true, false) + m.client.SendMessageEvent(evt.RoomID, event.EventMessage, &formattedResp) } } } diff --git a/gpt/gpt.go b/gpt/gpt.go deleted file mode 100644 index 56f2894..0000000 --- a/gpt/gpt.go +++ /dev/null @@ -1 +0,0 @@ -package gpt diff --git a/main.go b/main.go index 299b8aa..d12c523 100644 --- a/main.go +++ b/main.go @@ -4,33 +4,24 @@ import ( "os" "os/signal" - "ewintr.nl/matrix-bots/matrix" + "ewintr.nl/matrix-bots/bot" _ "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)) - matrixClient := matrix.New(matrix.Config{ + matrixClient := bot.New(bot.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"), + OpenAIKey: getParam("OPENAI_API_KEY", "no key"), }) - OpenaiAPIKey, ok := os.LookupEnv("OPENAI_API_KEY") - if !ok { - logger.Error("OPENAI_API_KEY is not set") - os.Exit(1) - } - - // Create new OpenAI client - openaiClient := openai.NewClient(OpenaiAPIKey) - if err := matrixClient.Init(); err != nil { logger.Error(err.Error()) os.Exit(1) @@ -38,7 +29,7 @@ func main() { go matrixClient.Run() matrixClient.AddEventHandler(matrixClient.InviteHandler()) - matrixClient.AddEventHandler(matrixClient.RespondHandler(openaiClient)) + matrixClient.AddEventHandler(matrixClient.RespondHandler()) done := make(chan os.Signal) signal.Notify(done, os.Interrupt)