llm: add & shortcut

config/config.go

@@ -179,6 +179,10 @@ func (c *Config) Set(key, value string) error {
 	return nil
 }
 
+func (c *Config) SetBool(key string, value bool) error {
+	return c.Set(key, fmt.Sprintf("%v", value))
+}
+
 func (c *Config) RefreshSecrets() error {
 	q := `select key, value from secrets`
 	var secrets []Secret

connectors/discord/discord.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"github.com/velour/catbase/plugins/talklikeapirate"
 	"net/http"
 	"strconv"
 	"strings"
@@ -33,6 +34,8 @@ type Discord struct {
 	cmdHandlers    map[string]CmdHandler
 
 	guildID string
+
+	Pirate *talklikeapirate.TalkLikeAPirateFilter
 }
 
 func New(config *config.Config) *Discord {
@@ -112,6 +115,14 @@ func (d Discord) Send(kind bot.Kind, args ...any) (string, error) {
 }
 
 func (d *Discord) sendMessage(channel, message string, meMessage bool, args ...any) (string, error) {
+	var err error
+	if d.Pirate != nil {
+		message, err = d.Pirate.Filter(message)
+		if err != nil {
+			log.Error().Err(err).Msg("could not pirate message")
+		}
+	}
+
 	if meMessage && !strings.HasPrefix(message, "_") && !strings.HasSuffix(message, "_") {
 		message = "_" + message + "_"
 	}
@@ -167,7 +178,6 @@ func (d *Discord) sendMessage(channel, message string, meMessage bool, args ...a
 	maxLen := 2000
 	chunkSize := maxLen - 100
 	var st *discordgo.Message
-	var err error
 	if len(data.Content) > maxLen {
 		tmp := data.Content
 		data.Content = tmp[:chunkSize]

main.go

@@ -7,6 +7,7 @@ package main
 import (
 	"flag"
 	"github.com/velour/catbase/plugins"
+	"github.com/velour/catbase/plugins/talklikeapirate"
 	"io"
 	"math/rand"
 	"os"
@@ -71,7 +72,9 @@ func main() {
 	case "slackapp":
 		client = slackapp.New(c)
 	case "discord":
-		client = discord.New(c)
+		d := discord.New(c)
+		d.Pirate = talklikeapirate.NewFilter(c)
+		client = d
 	default:
 		log.Fatal().Msgf("Unknown connection type: %s", c.Get("type", "UNSET"))
 	}

plugins/deadend/deadend.go

@@ -6,7 +6,7 @@ import (
 	"regexp"
 )
 
-const defaultMessage = "I don't know how to respond to that. If you'd like to ask an LLM, use the `llm` command."
+const defaultMessage = "I don't know how to respond to that. If you'd like to ask an LLM, use the `llm` command (or prefix your message with &)."
 
 type DeadEndPlugin struct {
 	b bot.Bot

plugins/llm/gemini.go

@@ -25,7 +25,7 @@ func (p *LLMPlugin) geminiConnect() error {
 }
 
 func (p *LLMPlugin) gemini(msg string) (chatEntry, error) {
-	model := p.geminiClient.GenerativeModel("gemini-1.5-flash")
+	model := p.geminiClient.GenerativeModel(p.c.Get("gemini.model", "gemini-1.5-flash"))
 	model.SetMaxOutputTokens(int32(p.c.GetInt("gemini.maxtokens", 100)))
 	model.SetTopP(float32(p.c.GetFloat64("gemini.topp", 0.95)))
 	model.SetTopK(int32(p.c.GetInt("gemini.topk", 20)))

plugins/llm/gpt.go

@@ -8,6 +8,7 @@ import (
 	"github.com/velour/catbase/bot"
 	"github.com/velour/catbase/config"
 	"regexp"
+	"strings"
 	"time"
 )
 
@@ -48,6 +49,12 @@ func (p *LLMPlugin) register() {
 			HelpText: "set the ChatGPT prompt",
 			Handler:  p.setPromptMessage,
 		},
+		{
+			Kind: bot.Message, IsCmd: false,
+			Regex:    regexp.MustCompile(`(?is)^&(?P<text>.*)`),
+			HelpText: "chat completion using first-available AI",
+			Handler:  p.geminiChatMessage,
+		},
 		{
 			Kind: bot.Message, IsCmd: true,
 			Regex:    regexp.MustCompile(`(?is)^llm (?P<text>.*)`),
@@ -55,7 +62,7 @@ func (p *LLMPlugin) register() {
 			Handler:  p.geminiChatMessage,
 		},
 		{
-			Kind: bot.Message, IsCmd: true,
+			Kind: bot.Message, IsCmd: false,
 			Regex:    regexp.MustCompile(`(?is)^gpt4 (?P<text>.*)`),
 			HelpText: "chat completion using OpenAI",
 			Handler:  p.gptMessage,
@@ -66,6 +73,11 @@ func (p *LLMPlugin) register() {
 			HelpText: "clear chat history",
 			Handler:  p.puke,
 		},
+		{
+			Kind: bot.Help, IsCmd: false,
+			Regex:   regexp.MustCompile(`.*`),
+			Handler: p.help,
+		},
 	}
 	p.b.RegisterTable(p, p.h)
 }
@@ -164,3 +176,16 @@ func (p *LLMPlugin) puke(r bot.Request) bool {
 	p.b.Send(r.Conn, bot.Message, r.Msg.Channel, resp)
 	return true
 }
+
+func (p *LLMPlugin) help(r bot.Request) bool {
+	out := "Talk like a pirate commands:\n"
+	for _, h := range p.h {
+		if h.HelpText == "" {
+			continue
+		}
+		out += fmt.Sprintf("```%s```\t%s", h.Regex.String(), h.HelpText)
+	}
+	out = strings.TrimSpace(out)
+	p.b.Send(r.Conn, bot.Message, r.Msg.Channel, out)
+	return true
+}

plugins/plugins.go

@@ -45,6 +45,7 @@ import (
 	"github.com/velour/catbase/plugins/sms"
 	"github.com/velour/catbase/plugins/stock"
 	"github.com/velour/catbase/plugins/talker"
+	"github.com/velour/catbase/plugins/talklikeapirate"
 	"github.com/velour/catbase/plugins/tappd"
 	"github.com/velour/catbase/plugins/tell"
 	"github.com/velour/catbase/plugins/tldr"
@@ -102,6 +103,7 @@ func Register(b bot.Bot) {
 	b.AddPlugin(talker.New(b))
 	b.AddPlugin(fact.New(b))
 	b.AddPlugin(llm.New(b))
+	b.AddPlugin(talklikeapirate.New(b))
 	// catches anything left, will always return true
 	b.AddPlugin(deadend.New(b))
 }

plugins/talklikeapirate/filter.go

@@ -0,0 +1,108 @@
+package talklikeapirate
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"github.com/google/generative-ai-go/genai"
+	"github.com/rs/zerolog/log"
+	"github.com/velour/catbase/bot"
+	"github.com/velour/catbase/config"
+	"google.golang.org/api/option"
+)
+
+// TalkLikeAPirateFilter reimplements the send function
+// with an AI intermediate.
+type TalkLikeAPirateFilter struct {
+	client *genai.Client
+	prompt string
+
+	b bot.Bot
+	c *config.Config
+}
+
+func NewFilter(c *config.Config) *TalkLikeAPirateFilter {
+	p := &TalkLikeAPirateFilter{
+		c: c,
+	}
+
+	return p
+}
+
+func (p *TalkLikeAPirateFilter) Filter(input string) (string, error) {
+	if !p.c.GetBool("talklikeapirate.enabled", false) {
+		return input, nil
+	}
+	if p.client == nil {
+		var err error
+		p.client, err = p.getClient()
+		if err != nil {
+			return input, err
+		}
+	}
+	model, err := p.GetModel()
+	if err != nil {
+		log.Error().Err(err).Send()
+		return input, err
+	}
+
+	res, err := model.GenerateContent(context.Background(), genai.Text(input))
+	if err != nil {
+		log.Error().Err(err).Send()
+		return input, err
+	}
+
+	if len(res.Candidates) == 0 {
+		err := errors.New("no candidates found")
+		log.Error().Err(err).Send()
+		return input, err
+	}
+
+	// Need to check here that we got an actual completion, not a
+	// warning about bad content. FinishReason exists on Completion.
+
+	completion := ""
+	for _, p := range res.Candidates[0].Content.Parts {
+		completion += fmt.Sprintf("%s", p)
+	}
+
+	return completion, nil
+}
+
+func (p *TalkLikeAPirateFilter) GetModel() (*genai.GenerativeModel, error) {
+	model := p.client.GenerativeModel(p.c.Get("gemini.model", "gemini-1.5-flash"))
+	model.SetMaxOutputTokens(int32(p.c.GetInt("gemini.maxtokens", 100)))
+	model.SetTopP(float32(p.c.GetFloat64("gemini.topp", 0.95)))
+	model.SetTopK(int32(p.c.GetInt("gemini.topk", 20)))
+	model.SetTemperature(float32(p.c.GetFloat64("gemini.temp", 0.9)))
+
+	model.SafetySettings = []*genai.SafetySetting{
+		{genai.HarmCategoryHarassment, genai.HarmBlockNone},
+		{genai.HarmCategoryHateSpeech, genai.HarmBlockNone},
+		{genai.HarmCategorySexuallyExplicit, genai.HarmBlockNone},
+		{genai.HarmCategoryDangerousContent, genai.HarmBlockNone},
+	}
+
+	if prompt := p.c.Get("talklikeapirate.systemprompt", ""); prompt != "" {
+		model.SystemInstruction = &genai.Content{
+			Parts: []genai.Part{genai.Text(prompt)},
+		}
+	} else {
+		return nil, errors.New("no system prompt selected")
+	}
+
+	return model, nil
+}
+
+func (p *TalkLikeAPirateFilter) getClient() (*genai.Client, error) {
+	ctx := context.Background()
+	key := p.c.Get("GEMINI_API_KEY", "")
+	if key == "" {
+		return nil, errors.New("missing GEMINI_API_KEY")
+	}
+	client, err := genai.NewClient(ctx, option.WithAPIKey(key))
+	if err != nil {
+		return nil, err
+	}
+	return client, nil
+}

plugins/talklikeapirate/plugin.go

@@ -0,0 +1,84 @@
+package talklikeapirate
+
+import (
+	"fmt"
+	"github.com/velour/catbase/bot"
+	"github.com/velour/catbase/config"
+	"regexp"
+	"strings"
+)
+
+// TalkLikeAPiratePlugin allows admin of the filter
+type TalkLikeAPiratePlugin struct {
+	b        bot.Bot
+	c        *config.Config
+	handlers bot.HandlerTable
+}
+
+func New(b bot.Bot) *TalkLikeAPiratePlugin {
+	p := &TalkLikeAPiratePlugin{
+		b: b,
+		c: b.Config(),
+	}
+
+	p.register()
+
+	return p
+}
+
+func (p *TalkLikeAPiratePlugin) register() {
+	p.handlers = bot.HandlerTable{
+		{
+			Kind: bot.Message, IsCmd: true,
+			Regex:    regexp.MustCompile(`^enable pirate$`),
+			HelpText: "Enable message filter",
+			Handler:  p.setEnabled(true),
+		},
+		{
+			Kind: bot.Message, IsCmd: true,
+			Regex:    regexp.MustCompile(`^disable pirate$`),
+			HelpText: "Disable message filter",
+			Handler:  p.setEnabled(false),
+		},
+		{
+			Kind: bot.Message, IsCmd: true,
+			Regex:    regexp.MustCompile(`^pirate-prompt:? (?P<text>.*)$`),
+			HelpText: "Set message filter prompt",
+			Handler:  p.setPrompt,
+		},
+		{
+			Kind: bot.Help, IsCmd: false,
+			Regex:   regexp.MustCompile(`.*`),
+			Handler: p.help,
+		},
+	}
+	p.b.RegisterTable(p, p.handlers)
+}
+
+func (p *TalkLikeAPiratePlugin) setEnabled(isEnabled bool) bot.ResponseHandler {
+	return func(r bot.Request) bool {
+		p.c.SetBool("talklikeapirate.enabled", isEnabled)
+		p.b.Send(r.Conn, bot.Message, r.Msg.Channel, fmt.Sprintf("I just set the message filter status to: %v", isEnabled))
+		return true
+	}
+}
+
+func (p *TalkLikeAPiratePlugin) setPrompt(r bot.Request) bool {
+	prompt := r.Values["text"]
+	p.c.Set("talklikeapirate.systemprompt", prompt)
+	p.b.Send(r.Conn, bot.Message, r.Msg.Channel, fmt.Sprintf("I set the message filter prompt to: %s", prompt))
+	return true
+}
+
+func (p *TalkLikeAPiratePlugin) help(r bot.Request) bool {
+	out := "Talk like a pirate commands:\n"
+	for _, h := range p.handlers {
+		if h.HelpText == "" {
+			continue
+		}
+		out += fmt.Sprintf("```%s```\t%s", h.Regex.String(), h.HelpText)
+	}
+	out = strings.TrimSpace(out)
+	p.b.Send(r.Conn, bot.Message, r.Msg.Channel, out)
+	return true
+}

plugins/tldr/tldr.go

@@ -129,7 +129,7 @@ func (p *TLDRPlugin) betterTLDR(r bot.Request) bool {
 	ch := r.Msg.Channel
 	c, err := p.getClient()
 	if err != nil {
-		p.b.Send(r.Conn, bot.Message, r.Msg.Channel, "Couldn't fetch an OpenAI client")
+		p.b.Send(r.Conn, bot.Message, r.Msg.Channel, "Couldn't fetch an AI client")
 		return true
 	}
 	promptConfig := p.c.Get(templateKey, defaultTemplate)
@@ -148,7 +148,7 @@ func (p *TLDRPlugin) betterTLDR(r bot.Request) bool {
 		backlog = str + backlog
 	}
 
-	model := c.GenerativeModel("gemini-1.5-flash")
+	model := c.GenerativeModel(p.c.Get("gemini.model", "gemini-1.5-flash"))
 	model.SystemInstruction = &genai.Content{
 		Parts: []genai.Part{genai.Text(prompt.String())},
 	}