Change to window memory.

jupyterlab · Apr 17, 2023 · e70fd6a · e70fd6a
1 parent 763e132
commit e70fd6a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py
@@ -1,6 +1,6 @@
 import asyncio
 import queue
-from langchain.memory import ConversationBufferMemory
+from langchain.memory import ConversationBufferWindowMemory
 from jupyter_ai.actors.default import DefaultActor 
 from jupyter_ai.actors.filesystem import FileSystemActor 
 from jupyter_ai.actors.index import DocumentIndexActor
@@ -125,7 +125,7 @@ def initialize_settings(self):
             log=self.log
         )
         memory_actor = MemoryActor.options(name=ACTOR_TYPE.MEMORY.value).remote(
-            memory=ConversationBufferMemory(return_messages=True),
+            memory=ConversationBufferWindowMemory(return_messages=True, k=2),
             log=self.log
         )
         self.settings['router'] = router