Added more memory management

6 files changed, 196 insertions, 5 deletions

diff --git a/agent.py b/agent.py
index 03abfe9..886cae0 100644
--- a/agent.py
+++ b/agent.py
@@ -4,7 +4,13 @@ from context import ContextBuilder
 from identity import AgentIdentity, UserProfile
 from llm import LLM
 from memory import WorkingMemory
-from tools import MemoryAddFactTool, MemoryAddNoteTool, ToolError, ToolRegistry
+from tools import (
+    MemoryAddFactTool,
+    MemoryAddNoteTool,
+    TaskStartTool,
+    ToolError,
+    ToolRegistry,
+)
 
 
 class Agent:
@@ -27,12 +33,11 @@ class Agent:
 
         self.tools.register(MemoryAddFactTool(self.working_memory))
         self.tools.register(MemoryAddNoteTool(self.working_memory))
+        self.tools.register(TaskStartTool(self.working_memory))
 
         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(
             {
diff --git a/agent.toml b/agent.toml
index 87189cd..7a3cfec 100644
--- a/agent.toml
+++ b/agent.toml
@@ -9,3 +9,61 @@ directness = true
 
 [quirks]
 favorite_color = "purple and orange"
+
+[behavior]
+task_management = """
+Determine whether the user is asking you to pursue an ongoing objective.
+
+Start a task with the task_start tool when:
+- the user explicitly asks you to work on a project, activity, or goal
+  over multiple interactions;
+- the user says they intend to focus on, continue working on, or spend
+  significant time developing something; or
+- the conversation clearly establishes an ongoing objective that you
+  are expected to continue working toward.
+
+Statements indicating an ongoing commitment, such as "over the next few
+days", "over the next few conversations", "I'm working on", "I'm building",
+or "I want to focus on", are strong signals that a task should be started.
+
+Do not start a task for:
+- casual conversation,
+- isolated questions,
+- one-off requests,
+- short-lived interactions,
+- or temporary interruptions to the current task.
+
+If a task is already active, continue working toward its objective unless
+the user clearly abandons it.
+
+When the user temporarily changes subjects or asks for a short-lived side
+request, do not replace the active task. Treat the request as a temporary
+interruption and preserve the current task.
+
+If the user explicitly abandons the current task and establishes a new
+ongoing objective, replace the current task with the new objective.
+
+Do not create a task merely because a conversation has multiple turns.
+"""
+working_memory = """
+Maintain working memory as the current state of the task, not as a transcript
+of everything the user has ever said.
+
+Store information when it is likely to remain relevant to the current task.
+
+When the user changes, corrects, or replaces previously stored information:
+- treat the new information as authoritative;
+- remove or update the stale information rather than keeping both;
+- do not preserve contradictory facts or notes as separate entries.
+
+Use facts for stable task-relevant information and decisions.
+Use notes for temporary information that may help with ongoing work.
+
+Do not store every conversational detail.
+Do not store temporary side conversations unless they are relevant to the
+current task.
+
+Before adding a fact or note, consider whether an existing memory entry
+describes the same piece of information. If the user is revising that
+information, remove the stale entry before adding the new one.
+"""
diff --git a/client.py b/client.py
index 438f1fb..d91d79d 100644
--- a/client.py
+++ b/client.py
@@ -5,7 +5,6 @@ from tools import (
     CalculatorTool,
     FileSystemListTool,
     FileSystemReadTool,
-    HelloTool,
     ToolRegistry,
     WeatherTool,
 )
@@ -18,7 +17,7 @@ agent_identity = load_agent_identity("agent.toml")
 user_profile = load_user_profile("user.toml")
 
 tools = ToolRegistry()
-tools.register(HelloTool())
+# tools.register(HelloTool())
 tools.register(CalculatorTool())
 tools.register(WeatherTool())
 tools.register(FileSystemListTool())
diff --git a/context.py b/context.py
index d702706..669e237 100644
--- a/context.py
+++ b/context.py
@@ -25,6 +25,9 @@ class ContextBuilder:
             f"Directness: {agent.values['directness']}.\n"
             f"Favorite color: {agent.quirks['favorite_color']}.\n"
             f"\n"
+            f"Task management:\n"
+            f"{agent.behavior['task_management']}"
+            f"\n"
             f"The user's name is {user.name}"
             f"\n"
             f"{self.build_working_memory(memory)}"
@@ -36,6 +39,9 @@ class ContextBuilder:
         if memory.objective:
             lines.append(f"Objective: {memory.objective}")
 
+        if memory.status:
+            lines.append(f"Status: {memory.status.value}")
+
         if memory.facts:
             lines.append("Facts:")
             lines.extend(f"- {fact}" for fact in memory.facts)
diff --git a/identity.py b/identity.py
index 312e3d7..ec1cc19 100644
--- a/identity.py
+++ b/identity.py
@@ -8,6 +8,7 @@ class AgentIdentity:
     personality: dict[str, str]
     values: dict[str, bool]
     quirks: dict[str, str]
+    behavior: dict[str, str]
 
 
 @dataclass
@@ -25,6 +26,7 @@ def load_agent_identity(path: str) -> AgentIdentity:
         personality=data["personality"],
         values=data["values"],
         quirks=data["quirks"],
+        behavior=data["behavior"],
     )
 
 
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: