telegram-ollama-reply-bot/main.go

91 lines
2.1 KiB
Go

package main
import (
"context"
"fmt"
"log/slog"
"os"
openai "github.com/sashabaranov/go-openai"
tg "github.com/mymmrac/telego"
tu "github.com/mymmrac/telego/telegoutil"
)
const (
ModelMistral = "mistral"
ModelMistralUncensored = "dolphin-mistral"
)
func main() {
ollamaToken := os.Getenv("OLLAMA_TOKEN")
ollamaBaseUrl := os.Getenv("OLLAMA_BASE_URL")
telegramToken := os.Getenv("TELEGRAM_TOKEN")
config := openai.DefaultConfig(ollamaToken)
config.BaseURL = ollamaBaseUrl
client := openai.NewClientWithConfig(config)
bot, err := tg.NewBot(telegramToken, tg.WithDefaultLogger(false, true))
if err != nil {
fmt.Println(err)
os.Exit(1)
}
// Get updates channel
updates, _ := bot.UpdatesViaLongPolling(nil)
// Stop reviving updates from update channel
defer bot.StopLongPolling()
// Loop through all updates when they came
for update := range updates {
// Check if update contains a message
if update.Message != nil {
slog.Info("Update with message received", update.Message.Chat, update.Message.From, update.Message.Text)
chatID := tu.ID(update.Message.Chat.ID)
req := openai.ChatCompletionRequest{
Model: ModelMistralUncensored,
Messages: []openai.ChatCompletionMessage{
{
Role: openai.ChatMessageRoleSystem,
Content: "You're a bot in the Telegram chat. You are replying to questions directed to you.",
},
},
}
req.Messages = append(req.Messages, openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleUser,
Content: update.Message.Text,
})
resp, err := client.CreateChatCompletion(context.Background(), req)
if err != nil {
slog.Error("ChatCompletion error", err)
continue
}
slog.Info("Got completion. Going to send.", resp.Choices[0])
message := tu.Message(
chatID,
resp.Choices[0].Message.Content,
)
message = message.WithReplyParameters(&tg.ReplyParameters{MessageID: update.Message.MessageID})
message = message.WithParseMode("Markdown")
_, err = bot.SendMessage(message)
if err != nil {
slog.Error("Can't send reply message", err)
}
//req.Messages = append(req.Messages, resp.Choices[0].Message)
}
}
}