Added more memrory methods

1 files changed, 72 insertions, 0 deletions

diff --git a/memory.py b/memory.py
index 9269f94..6ba3368 100644
--- a/memory.py
+++ b/memory.py
@@ -1,4 +1,11 @@
 from dataclasses import dataclass, field
+from enum import Enum
+
+
+class TaskStatus(Enum):
+    ACTIVE = "active"
+    PAUSED = "paused"
+    COMPLETED = "completed"
 
 
 @dataclass
@@ -6,3 +13,68 @@ class WorkingMemory:
     objective: str | None = None
     facts: list[str] = field(default_factory=list)
     notes: list[str] = field(default_factory=list)
+    status: TaskStatus | None = None
+
+    def __str__(self) -> str:
+        parts: list[str] = []
+
+        if self.objective:
+            parts.append(f"Objective: {self.objective}")
+
+        if self.status:
+            parts.append(f"Status: {self.status.value}")
+
+        if self.facts:
+            parts.append("Facts:")
+            parts.extend(f"- {fact}" for fact in self.facts)
+
+        if self.notes:
+            parts.append("Notes:")
+            parts.extend(f"- {note}" for note in self.notes)
+
+        return "\n".join(parts) if parts else "(empty)"
+
+    def set_objective(self, objective: str) -> None:
+        self.objective = objective
+
+    def add_fact(self, fact: str) -> None:
+        self.facts.append(fact)
+
+    def remove_fact(self, fact: str) -> None:
+        if fact in self.facts:
+            self.facts.remove(fact)
+
+    def add_note(self, note: str) -> None:
+        self.notes.append(note)
+
+    def remove_note(self, note: str) -> None:
+        if note in self.notes:
+            self.notes.remove(note)
+
+    def set_status(self, status: TaskStatus):
+        self.status = status
+
+    def clear(self) -> None:
+        self.objective = None
+        self.status = None
+        self.facts.clear()
+        self.notes.clear()
+
+    def start_task(self, objective: str) -> None:
+        self.clear()
+        self.objective = objective
+        self.set_status(TaskStatus.ACTIVE)
+
+    def pause_task(self) -> None:
+        if self.status == TaskStatus.ACTIVE:
+            self.set_status(TaskStatus.PAUSED)
+
+    def resume_task(self) -> None:
+        if self.status == TaskStatus.PAUSED:
+            self.set_status(TaskStatus.ACTIVE)
+
+    def complete_task(self) -> None:
+        self.set_status(TaskStatus.COMPLETED)
+
+    def end_task(self) -> None:
+        self.clear()