feat(agent): shared state, allow to track conversations globally (#148)

* feat(agent): shared state, allow to track conversations globally

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Cleanup

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* track conversations initiated by the bot

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

core/conversations/conversations_suite_test.go

@@ -0,0 +1,13 @@
+package conversations_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestConversations(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Conversations test suite")
+}

core/conversations/conversationstracker.go

@@ -0,0 +1,84 @@
+package conversations
+
+import (
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/mudler/LocalAGI/pkg/xlog"
+	"github.com/sashabaranov/go-openai"
+)
+
+type TrackerKey interface{ ~int | ~int64 | ~string }
+
+type ConversationTracker[K TrackerKey] struct {
+	convMutex           sync.Mutex
+	currentconversation map[K][]openai.ChatCompletionMessage
+	lastMessageTime     map[K]time.Time
+	lastMessageDuration time.Duration
+}
+
+func NewConversationTracker[K TrackerKey](lastMessageDuration time.Duration) *ConversationTracker[K] {
+	return &ConversationTracker[K]{
+		lastMessageDuration: lastMessageDuration,
+		currentconversation: map[K][]openai.ChatCompletionMessage{},
+		lastMessageTime:     map[K]time.Time{},
+	}
+}
+
+func (c *ConversationTracker[K]) GetConversation(key K) []openai.ChatCompletionMessage {
+	// Lock the conversation mutex to update the conversation history
+	c.convMutex.Lock()
+	defer c.convMutex.Unlock()
+
+	// Clear up the conversation if the last message was sent more than lastMessageDuration ago
+	currentConv := []openai.ChatCompletionMessage{}
+	lastMessageTime := c.lastMessageTime[key]
+	if lastMessageTime.IsZero() {
+		lastMessageTime = time.Now()
+	}
+	if lastMessageTime.Add(c.lastMessageDuration).Before(time.Now()) {
+		currentConv = []openai.ChatCompletionMessage{}
+		c.lastMessageTime[key] = time.Now()
+		xlog.Debug("Conversation history does not exist for", "key", fmt.Sprintf("%v", key))
+	} else {
+		xlog.Debug("Conversation history exists for", "key", fmt.Sprintf("%v", key))
+		currentConv = append(currentConv, c.currentconversation[key]...)
+	}
+
+	// cleanup other conversations if older
+	for k := range c.currentconversation {
+		lastMessage, exists := c.lastMessageTime[k]
+		if !exists {
+			delete(c.currentconversation, k)
+			delete(c.lastMessageTime, k)
+			continue
+		}
+		if lastMessage.Add(c.lastMessageDuration).Before(time.Now()) {
+			xlog.Debug("Cleaning up conversation for", k)
+			delete(c.currentconversation, k)
+			delete(c.lastMessageTime, k)
+		}
+	}
+
+	return currentConv
+
+}
+
+func (c *ConversationTracker[K]) AddMessage(key K, message openai.ChatCompletionMessage) {
+	// Lock the conversation mutex to update the conversation history
+	c.convMutex.Lock()
+	defer c.convMutex.Unlock()
+
+	c.currentconversation[key] = append(c.currentconversation[key], message)
+	c.lastMessageTime[key] = time.Now()
+}
+
+func (c *ConversationTracker[K]) SetConversation(key K, messages []openai.ChatCompletionMessage) {
+	// Lock the conversation mutex to update the conversation history
+	c.convMutex.Lock()
+	defer c.convMutex.Unlock()
+
+	c.currentconversation[key] = messages
+	c.lastMessageTime[key] = time.Now()
+}

core/conversations/conversationstracker_test.go

@@ -0,0 +1,111 @@
+package conversations_test
+
+import (
+	"time"
+
+	"github.com/mudler/LocalAGI/core/conversations"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	"github.com/sashabaranov/go-openai"
+)
+
+var _ = Describe("ConversationTracker", func() {
+	var (
+		tracker  *conversations.ConversationTracker[string]
+		duration time.Duration
+	)
+
+	BeforeEach(func() {
+		duration = 1 * time.Second
+		tracker = conversations.NewConversationTracker[string](duration)
+	})
+
+	It("should initialize with empty conversations", func() {
+		Expect(tracker.GetConversation("test")).To(BeEmpty())
+	})
+
+	It("should add a message and retrieve it", func() {
+		message := openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: "Hello",
+		}
+		tracker.AddMessage("test", message)
+		conv := tracker.GetConversation("test")
+		Expect(conv).To(HaveLen(1))
+		Expect(conv[0]).To(Equal(message))
+	})
+
+	It("should clear the conversation after the duration", func() {
+		message := openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: "Hello",
+		}
+		tracker.AddMessage("test", message)
+		time.Sleep(2 * time.Second)
+		conv := tracker.GetConversation("test")
+		Expect(conv).To(BeEmpty())
+	})
+
+	It("should keep the conversation within the duration", func() {
+		message := openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: "Hello",
+		}
+		tracker.AddMessage("test", message)
+		time.Sleep(500 * time.Millisecond) // Half the duration
+		conv := tracker.GetConversation("test")
+		Expect(conv).To(HaveLen(1))
+		Expect(conv[0]).To(Equal(message))
+	})
+
+	It("should handle multiple keys and clear old conversations", func() {
+		message1 := openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: "Hello 1",
+		}
+		message2 := openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: "Hello 2",
+		}
+
+		tracker.AddMessage("key1", message1)
+		tracker.AddMessage("key2", message2)
+
+		time.Sleep(2 * time.Second)
+
+		conv1 := tracker.GetConversation("key1")
+		conv2 := tracker.GetConversation("key2")
+
+		Expect(conv1).To(BeEmpty())
+		Expect(conv2).To(BeEmpty())
+	})
+
+	It("should handle different key types", func() {
+		trackerInt := conversations.NewConversationTracker[int](duration)
+		trackerInt64 := conversations.NewConversationTracker[int64](duration)
+
+		message := openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: "Hello",
+		}
+
+		trackerInt.AddMessage(1, message)
+		trackerInt64.AddMessage(int64(1), message)
+
+		Expect(trackerInt.GetConversation(1)).To(HaveLen(1))
+		Expect(trackerInt64.GetConversation(int64(1))).To(HaveLen(1))
+	})
+
+	It("should cleanup other conversations if older", func() {
+		message := openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: "Hello",
+		}
+		tracker.AddMessage("key1", message)
+		tracker.AddMessage("key2", message)
+		time.Sleep(2 * time.Second)
+		tracker.GetConversation("key3")
+		Expect(tracker.GetConversation("key1")).To(BeEmpty())
+		Expect(tracker.GetConversation("key2")).To(BeEmpty())
+	})
+})