Evict memory stream (#5)

Author: fogfish

examples/chain/chain.go

@@ -65,7 +65,7 @@ func NewAgentB(llm chatter.Chatter) *AgentB {
 	agt.registry.Register(command.Return())
 
 	agt.Automata = agent.NewAutomata(llm,
-		memory.NewStream(`
+		memory.NewStream(memory.INFINITE, `
 			You are automomous agent who uses tools to perform required tasks.
 			You are using and remember context from earlier chat history to execute the task.
 		`),

examples/rainbow/rainbow.go

@@ -97,7 +97,7 @@ func main() {
 
 		// Configures memory for the agent. Typically, memory retains all of
 		// the agent's observations. Here, we use a stream memory that holds all observations.
-		memory.NewStream("You are agent who remembers and uses earlier chat history."),
+		memory.NewStream(memory.INFINITE, "You are agent who remembers and uses earlier chat history."),
 
 		// Configures the reasoner, which determines the agent's next actions and prompts.
 		// Here, we use custom (app specific) reasoner. The agent is restricted to execute

examples/script/script.go

@@ -95,7 +95,7 @@ func main() {
 
 		// Configures memory for the agent. Typically, memory retains all of
 		// the agent's observations. Here, we use a stream memory that holds all observations.
-		memory.NewStream(`
+		memory.NewStream(memory.INFINITE, `
 			You are automomous agent who uses tools to perform required tasks.
 			You are using and remember context from earlier chat history to execute the task.
 		`),

memory/stream.go

@@ -17,20 +17,24 @@ import (
 	"github.com/kshard/thinker"
 )
 
+const INFINITE = -1
+
 // The stream memory retains all of the agent's observations in the time ordered sequence.
 type Stream struct {
 	mu      sync.Mutex
 	heap    map[guid.K]*thinker.Observation
 	commits []guid.K
 	stratum chatter.Stratum
+	cap     int
 }
 
 // Creates new stream memory that retains all of the agent's observations.
-func NewStream(stratum chatter.Stratum) *Stream {
+func NewStream(cap int, stratum chatter.Stratum) *Stream {
 	return &Stream{
 		heap:    make(map[guid.K]*thinker.Observation),
 		commits: make([]guid.K, 0),
 		stratum: stratum,
+		cap:     cap,
 	}
 }
 
@@ -41,6 +45,12 @@ func (s *Stream) Commit(e *thinker.Observation) {
 
 	s.heap[e.Created] = e
 	s.commits = append(s.commits, e.Created)
+
+	if s.cap > 0 && len(s.commits) > s.cap {
+		head := s.commits[0]
+		s.commits = s.commits[1:]
+		delete(s.heap, head)
+	}
 }
 
 // Builds the context window for LLM using incoming prompt.

memory/stream_test.go

@@ -17,7 +17,7 @@ import (
 )
 
 func TestStream(t *testing.T) {
-	s := NewStream("role.")
+	s := NewStream(2, "role.")
 
 	t.Run("Commit", func(t *testing.T) {
 		var prompt chatter.Prompt
@@ -45,4 +45,25 @@ func TestStream(t *testing.T) {
 			it.Seq(seq).Equal("role.", "prompt.", "reply.", "ask."),
 		)
 	})
+
+	t.Run("Evict", func(t *testing.T) {
+		a := thinker.NewObservation(chatter.Prompt{Task: "0."}, chatter.Reply{Text: "0."})
+		s.Commit(a)
+
+		b := thinker.NewObservation(chatter.Prompt{Task: "a."}, chatter.Reply{Text: "a."})
+		s.Commit(b)
+
+		c := thinker.NewObservation(chatter.Prompt{Task: "b."}, chatter.Reply{Text: "b."})
+		s.Commit(c)
+
+		seq := make([]string, 0)
+		for _, x := range s.Context(chatter.Prompt{Task: "c."}) {
+			seq = append(seq, x.String())
+		}
+
+		it.Then(t).Should(
+			it.Seq(seq).Equal("role.", "a.", "a.", "b.", "b.", "c."),
+		)
+	})
+
 }