Added more memory management

1 files changed, 121 insertions, 0 deletions

diff --git a/tools.py b/tools.py
index c40dcd0..54c837e 100644
--- a/tools.py
+++ b/tools.py
@@ -61,6 +61,55 @@ class ToolRegistry:
         ]
 
 
+class TaskStartTool(Tool):
+    def __init__(self, memory: WorkingMemory) -> None:
+        self.memory = memory
+
+    @property
+    def name(self) -> str:
+        return "task_start"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Record the user's ongoing project or goal as the current task. "
+            "Use this when the user intends to work on something across multiple "
+            "turns or over a period of time. Do not use it for one-off requests "
+            "or temporary side conversations."
+        )
+
+    @property
+    def schema(self) -> dict[str, object]:
+        return {
+            "type": "object",
+            "properties": {
+                "objective": {
+                    "type": "string",
+                    "description": (
+                        "A concise description of the ongoing objective the user wants "
+                        "to pursue. For example: 'Build an image tagging API' or "
+                        "'Design and implement a Rust CLI that converts directory trees "
+                        "into static HTML.'"
+                    ),
+                },
+            },
+            "required": ["objective"],
+        }
+
+    def execute(self, arguments: dict[str, object]) -> ToolResult:
+        objective = arguments.get("objective")
+
+        if not isinstance(objective, str):
+            raise ToolError("objective must be a string")
+
+        if not objective.strip():
+            raise ToolError("objective must not be empty")
+
+        self.memory.start_task(objective)
+
+        return ToolResult(content=f"Started task: {objective}")
+
+
 class MemoryAddFactTool(Tool):
     def __init__(self, memory: WorkingMemory) -> None:
         self.memory = memory
@@ -133,6 +182,78 @@ class MemoryAddNoteTool(Tool):
         return ToolResult(content=f"Added working-memory note: {note}")
 
 
+class MemoryRemoveFactTool(Tool):
+    def __init__(self, memory: WorkingMemory) -> None:
+        self.memory = memory
+
+    @property
+    def name(self) -> str:
+        return "memory_remove_fact"
+
+    @property
+    def description(self) -> str:
+        return "Remove a fact from the working memory when it is no longer true or relevant."
+
+    @property
+    def schema(self) -> dict[str, object]:
+        return {
+            "type": "object",
+            "properties": {
+                "fact": {
+                    "type": "string",
+                    "description": "The exact fact to remove from working memory.",
+                },
+            },
+            "required": ["fact"],
+        }
+
+    def execute(self, arguments: dict[str, object]) -> ToolResult:
+        fact = arguments.get("fact")
+
+        if not isinstance(fact, str):
+            raise ToolError("fact must be a string")
+
+        self.memory.add_fact(fact)
+
+        return ToolResult(content=f"Removed fact: {fact}")
+
+
+class MemoryRemoveNoteTool(Tool):
+    def __init__(self, memory: WorkingMemory) -> None:
+        self.memory = memory
+
+    @property
+    def name(self) -> str:
+        return "memory_remove_note"
+
+    @property
+    def description(self) -> str:
+        return "Remove a note from the working memory when it is no longer true or relevant."
+
+    @property
+    def schema(self) -> dict[str, object]:
+        return {
+            "type": "object",
+            "properties": {
+                "note": {
+                    "type": "string",
+                    "description": "The exact note to remove from working memory.",
+                },
+            },
+            "required": ["note"],
+        }
+
+    def execute(self, arguments: dict[str, object]) -> ToolResult:
+        note = arguments.get("note")
+
+        if not isinstance(note, str):
+            raise ToolError("note must be a string")
+
+        self.memory.add_note(note)
+
+        return ToolResult(content=f"Removed note: {note}")
+
+
 class HelloTool(Tool):
     @property
     def name(self) -> str: