From 0a12e796d0bae4001fb8cf83ff815f9c6844c771 Mon Sep 17 00:00:00 2001
From: Chris Sexton <3216719+chrissexton@users.noreply.github.com>
Date: Fri, 27 Sep 2024 17:12:47 -0400
Subject: [PATCH] llm: timeout old messages

---
 plugins/llm/gemini.go |  2 +-
 plugins/llm/gpt.go    | 39 +++++++++++++++++++++++++++++++++++----
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/plugins/llm/gemini.go b/plugins/llm/gemini.go
index a2ee7fe..1974ae2 100644
--- a/plugins/llm/gemini.go
+++ b/plugins/llm/gemini.go
@@ -49,7 +49,7 @@ func (p *LLMPlugin) gemini(msg string) (chatEntry, error) {
 	ctx := context.Background()
 
 	cs.History = []*genai.Content{}
-	for _, h := range p.chatHistory {
+	for _, h := range p.getChatHistory() {
 		cs.History = append(cs.History, &genai.Content{
 			Parts: []genai.Part{
 				genai.Text(h.Content),
diff --git a/plugins/llm/gpt.go b/plugins/llm/gpt.go
index 1626c2b..39fd479 100644
--- a/plugins/llm/gpt.go
+++ b/plugins/llm/gpt.go
@@ -8,6 +8,7 @@ import (
 	"github.com/velour/catbase/bot"
 	"github.com/velour/catbase/config"
 	"regexp"
+	"time"
 )
 
 const gpt3URL = "https://api.openai.com/v1/engines/%s/completions"
@@ -25,8 +26,9 @@ type LLMPlugin struct {
 }
 
 type chatEntry struct {
-	Role    string `json:"role"`
-	Content string `json:"content"`
+	Role    string    `json:"role"`
+	Content string    `json:"content"`
+	TS      time.Time `json:"ts"`
 }
 
 func New(b bot.Bot) *LLMPlugin {
@@ -74,6 +76,35 @@ func (p *LLMPlugin) setPromptMessage(r bot.Request) bool {
 	return true
 }
 
+const defaultDuration = 15 * time.Minute
+
+func (p *LLMPlugin) getChatHistory() []chatEntry {
+	horizonTxt := p.c.Get("gemini.horizon", defaultDuration.String())
+	dur, err := time.ParseDuration(horizonTxt)
+	if err != nil {
+		dur = defaultDuration
+	}
+	output := []chatEntry{}
+	for _, e := range p.chatHistory {
+		if e.TS.After(time.Now().Add(-dur)) {
+			output = append(output, e)
+		}
+	}
+	return output
+}
+
+func (p *LLMPlugin) addChatHistoryUser(content string) {
+	p.addChatHistory(chatEntry{
+		Role:    "user",
+		Content: content,
+	})
+}
+
+func (p *LLMPlugin) addChatHistory(content chatEntry) {
+	content.TS = time.Now()
+	p.chatHistory = append(p.chatHistory, content)
+}
+
 func (p *LLMPlugin) geminiChatMessage(r bot.Request) bool {
 	if p.geminiClient == nil && p.geminiConnect() != nil {
 		log.Error().Msgf("Could not connect to Gemini")
@@ -85,8 +116,8 @@ func (p *LLMPlugin) geminiChatMessage(r bot.Request) bool {
 		p.b.Send(r.Conn, bot.Message, r.Msg.Channel, fmt.Sprintf("Problem with Gemini: %s", err))
 		return true
 	}
-	p.chatHistory = append(p.chatHistory, chatEntry{"User", r.Values["text"]})
-	p.chatHistory = append(p.chatHistory, chatResp)
+	p.addChatHistoryUser(r.Values["text"])
+	p.addChatHistory(chatResp)
 	p.b.Send(r.Conn, bot.Message, r.Msg.Channel, chatResp.Content)
 	log.Info().Msgf("Successfully used Gemini")
 	return true