Added more memrory methods

1 files changed, 20 insertions, 6 deletions

diff --git a/agent.py b/agent.py
index 41c7a7b..9dfd4ca 100644
--- a/agent.py
+++ b/agent.py
@@ -23,16 +23,14 @@ class Agent:
             user_profile,
         )
 
-        self.working_memory = WorkingMemory(
-            objective="Testing working memory",
-            facts=["The user is testing the agent's working memory system."],
-            notes=["This is temporary test data."],
-        )
+        self.working_memory = WorkingMemory()
 
         self.messages: list[dict[str, object]] = []
 
     def chat(self, user_input: str) -> str:
-
+        if self.working_memory.objective is None:
+            self.working_memory.start_task(user_input)
+        print(f"\n[memory: {self.working_memory}]")
         self.messages.append(
             {
                 "role": "user",
@@ -90,6 +88,7 @@ class Agent:
                 else:
                     try:
                         result = tool.execute(call.arguments).content
+                        self.working_memory.add_note(f"{call.name} returned: {result}")
                     except Exception as error:
                         result = f"Tool error: {error}"
 
@@ -101,3 +100,18 @@ class Agent:
                     }
                 )
         raise RuntimeError("Tool call limit exceeded")
+
+    def end_task(self) -> None:
+        self.working_memory.end_task()
+
+    def start_task(self, objective: str) -> None:
+        self.working_memory.start_task(objective)
+
+    def pause_task(self) -> None:
+        self.working_memory.pause_task()
+
+    def resume_task(self) -> None:
+        self.working_memory.resume_task()
+
+    def complete_task(self) -> None:
+        self.working_memory.complete_task()