Added chatgpt support

3 files changed, 50 insertions, 1 deletions

diff --git a/go.mod b/go.mod
index 07b5404..0d90b86 100644
--- a/go.mod
+++ b/go.mod
@@ -7,6 +7,7 @@ require (
 	github.com/google/generative-ai-go v0.7.0
 	github.com/lrstanley/girc v0.0.0-20240125042120-9add3166e52e
 	github.com/pelletier/go-toml/v2 v2.1.1
+	github.com/sashabaranov/go-openai v1.19.3
 	google.golang.org/api v0.163.0
 )
 
diff --git a/go.sum b/go.sum
index 4da9549..ef6c962 100644
--- a/go.sum
+++ b/go.sum
@@ -82,6 +82,8 @@ github.com/pelletier/go-toml/v2 v2.1.1/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdU
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
+github.com/sashabaranov/go-openai v1.19.3 h1:xJvkU8Tye6MOKLaoqjh7qXYwKiEYGtlmp06cb8179yo=
+github.com/sashabaranov/go-openai v1.19.3/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
diff --git a/main.go b/main.go
index 141625b..69d909f 100644
--- a/main.go
+++ b/main.go
@@ -17,6 +17,7 @@ import (
 	"github.com/google/generative-ai-go/genai"
 	"github.com/lrstanley/girc"
 	"github.com/pelletier/go-toml/v2"
+	openai "github.com/sashabaranov/go-openai"
 	"google.golang.org/api/option"
 )
 
@@ -221,7 +222,52 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client) {
 					return
 				}
 
-				fmt.Println(writer.String())
+				log.Println(writer.String())
+				client.Cmd.ReplyTo(event, girc.Fmt("\033[0m"+writer.String()))
+			}
+		})
+	} else if appConfig.Provider == "chatgpt" {
+		irc.Handlers.AddBg(girc.PRIVMSG, func(client *girc.Client, event girc.Event) {
+			if strings.HasPrefix(event.Last(), appConfig.IrcNick+": ") {
+				prompt := strings.TrimPrefix(event.Last(), appConfig.IrcNick+": ")
+				log.Println(prompt)
+
+				ctx, cancel := context.WithTimeout(context.Background(), time.Duration(appConfig.RequestTimeout)*time.Second)
+				defer cancel()
+
+				gptClient := openai.NewClient(appConfig.Apikey)
+
+				messages := make([]openai.ChatCompletionMessage, 0)
+
+				messages = append(messages, openai.ChatCompletionMessage{
+					Role:    "system",
+					Content: prompt,
+				})
+
+				resp, err := gptClient.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
+					Model:    appConfig.Model,
+					Messages: messages,
+				})
+				if err != nil {
+					client.Cmd.ReplyTo(event, girc.Fmt(fmt.Sprintf("error: %s", err.Error())))
+
+					return
+				}
+
+				var writer bytes.Buffer
+				err = quick.Highlight(
+					&writer,
+					resp.Choices[0].Message.Content,
+					"markdown",
+					appConfig.ChromaFormatter,
+					appConfig.ChromaStyle)
+				if err != nil {
+					client.Cmd.ReplyTo(event, girc.Fmt(fmt.Sprintf("error: %s", err.Error())))
+
+					return
+				}
+
+				log.Println(writer.String())
 				client.Cmd.ReplyTo(event, girc.Fmt("\033[0m"+writer.String()))
 			}
 		})