From 0dfefe297f0dbcdcf45d9ac2ff6d62a398572110 Mon Sep 17 00:00:00 2001
From: Chris Sexton <3216719+chrissexton@users.noreply.github.com>
Date: Fri, 27 Sep 2024 11:40:48 -0400
Subject: [PATCH] llm: add system prompt and fix vars

---
 plugins/llm/gemini.go | 14 ++++++++++----
 plugins/llm/gpt.go    |  6 +-----
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/plugins/llm/gemini.go b/plugins/llm/gemini.go
index 722d1b7..a2ee7fe 100644
--- a/plugins/llm/gemini.go
+++ b/plugins/llm/gemini.go
@@ -26,10 +26,10 @@ func (p *LLMPlugin) geminiConnect() error {
 
 func (p *LLMPlugin) gemini(msg string) (chatEntry, error) {
 	model := p.geminiClient.GenerativeModel("gemini-1.5-flash")
-	model.SetMaxOutputTokens(int32(p.c.GetInt("GEMINI_MAX_TOKENS", 100)))
-	model.SetTopP(float32(p.c.GetFloat64("GEMINI_TOP_P", 0.95)))
-	model.SetTopK(int32(p.c.GetInt("GEMINI_TOP_K", 20)))
-	model.SetTemperature(float32(p.c.GetFloat64("GEMINI_TEMP", 0.9)))
+	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},
@@ -38,6 +38,12 @@ func (p *LLMPlugin) gemini(msg string) (chatEntry, error) {
 		{genai.HarmCategoryDangerousContent, genai.HarmBlockNone},
 	}
 
+	if prompt := p.c.Get("gemini.systemprompt", ""); prompt != "" {
+		model.SystemInstruction = &genai.Content{
+			Parts: []genai.Part{genai.Text(prompt)},
+		}
+	}
+
 	cs := model.StartChat()
 
 	ctx := context.Background()
diff --git a/plugins/llm/gpt.go b/plugins/llm/gpt.go
index 4094f1c..14aed20 100644
--- a/plugins/llm/gpt.go
+++ b/plugins/llm/gpt.go
@@ -69,11 +69,7 @@ func (p *LLMPlugin) register() {
 }
 
 func (p *LLMPlugin) setPromptMessage(r bot.Request) bool {
-	prompt := r.Values["text"]
-	if err := p.setPrompt(prompt); err != nil {
-		resp := fmt.Sprintf("Error: %s", err)
-		p.b.Send(r.Conn, bot.Message, r.Msg.Channel, resp)
-	}
+	p.c.Set("gemini.systemprompt", r.Values["text"])
 	p.b.Send(r.Conn, bot.Message, r.Msg.Channel, fmt.Sprintf(`Okay. I set the prompt to: "%s"`, prompt))
 	return true
 }