Added memory api

5 files changed, 112 insertions, 29 deletions

diff --git a/agent.py b/agent.py
index 9dfd4ca..03abfe9 100644
--- a/agent.py
+++ b/agent.py
@@ -3,8 +3,8 @@ import json
 from context import ContextBuilder
 from identity import AgentIdentity, UserProfile
 from llm import LLM
-from tools import ToolRegistry
 from memory import WorkingMemory
+from tools import MemoryAddFactTool, MemoryAddNoteTool, ToolError, ToolRegistry
 
 
 class Agent:
@@ -25,6 +25,9 @@ class Agent:
 
         self.working_memory = WorkingMemory()
 
+        self.tools.register(MemoryAddFactTool(self.working_memory))
+        self.tools.register(MemoryAddNoteTool(self.working_memory))
+
         self.messages: list[dict[str, object]] = []
 
     def chat(self, user_input: str) -> str:
@@ -88,8 +91,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:
+                    except ToolError as error:
                         result = f"Tool error: {error}"
 
                 self.messages.append(
diff --git a/identity.py b/identity.py
index cbcd6ad..312e3d7 100644
--- a/identity.py
+++ b/identity.py
@@ -1,5 +1,5 @@
-from dataclasses import dataclass
 import tomllib
+from dataclasses import dataclass
 
 
 @dataclass
diff --git a/llm.py b/llm.py
index 4dde6bd..afd890c 100644
--- a/llm.py
+++ b/llm.py
@@ -1,6 +1,5 @@
-from dataclasses import dataclass
 import json
-from collections.abc import Iterator
+from dataclasses import dataclass
 
 import requests
 
@@ -68,7 +67,7 @@ class LLM:
 
 
 if __name__ == "__main__":
-    from tools import ToolRegistry, HelloTool
+    from tools import HelloTool, ToolRegistry
 
     registry = ToolRegistry()
     registry.register(HelloTool())
diff --git a/pyproject.toml b/pyproject.toml
index b98d3dc..b32d935 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,3 +9,6 @@ dependencies = [
 
 [tool.basedpyright]
 typeCheckingMode = "standard"
+
+[tool.ruff.lint]
+extend-select = ["I"]
diff --git a/tools.py b/tools.py
index 2ae8857..c40dcd0 100644
--- a/tools.py
+++ b/tools.py
@@ -1,8 +1,15 @@
-from datetime import datetime, timedelta, timezone, UTC
-from pathlib import Path
-import requests
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta, timezone
+from pathlib import Path
+
+import requests
+
+from memory import WorkingMemory
+
+
+class ToolError(Exception):
+    pass
 
 
 @dataclass
@@ -54,6 +61,78 @@ class ToolRegistry:
         ]
 
 
+class MemoryAddFactTool(Tool):
+    def __init__(self, memory: WorkingMemory) -> None:
+        self.memory = memory
+
+    @property
+    def name(self) -> str:
+        return "memory_add_fact"
+
+    @property
+    def description(self) -> str:
+        return "Add a useful fact to the current working memory."
+
+    @property
+    def schema(self) -> dict[str, object]:
+        return {
+            "type": "object",
+            "properties": {
+                "fact": {
+                    "type": "string",
+                    "description": "A fact worth remembering during the current task.",
+                },
+            },
+            "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"Added working-memory fact: {fact}")
+
+
+class MemoryAddNoteTool(Tool):
+    def __init__(self, memory: WorkingMemory) -> None:
+        self.memory = memory
+
+    @property
+    def name(self) -> str:
+        return "memory_add_note"
+
+    @property
+    def description(self) -> str:
+        return "Add a useful note to the current working memory."
+
+    @property
+    def schema(self) -> dict[str, object]:
+        return {
+            "type": "object",
+            "properties": {
+                "note": {
+                    "type": "string",
+                    "description": "A useful observation, reminder, or note for the current task.",
+                },
+            },
+            "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"Added working-memory note: {note}")
+
+
 class HelloTool(Tool):
     @property
     def name(self) -> str:
@@ -118,13 +197,13 @@ class CalculatorTool(Tool):
         operation = arguments["operation"]
 
         if not isinstance(a, (int, float)):
-            raise TypeError("a must be a number")
+            raise ToolError("a must be a number")
 
         if not isinstance(b, (int, float)):
-            raise TypeError("b must be a number")
+            raise ToolError("b must be a number")
 
         if not isinstance(operation, str):
-            raise TypeError("operation must be a string")
+            raise ToolError("operation must be a string")
 
         if operation == "add":
             result = a + b
@@ -134,10 +213,10 @@ class CalculatorTool(Tool):
             result = a * b
         elif operation == "divide":
             if b == 0:
-                raise ValueError("cannot divide by zero")
+                raise ToolError("cannot divide by zero")
             result = a / b
         else:
-            raise ValueError(f"unknown operation: {operation}")
+            raise ToolError(f"unknown operation: {operation}")
 
         return ToolResult(content=str(result))
 
@@ -168,7 +247,7 @@ class WeatherTool(Tool):
         location = arguments.get("location")
 
         if not isinstance(location, str):
-            raise TypeError("location must be a string")
+            raise ToolError("location must be a string")
 
         geocoding_response = requests.get(
             "https://geocoding-api.open-meteo.com/v1/search",
@@ -186,7 +265,7 @@ class WeatherTool(Tool):
         results = geocoding_data.get("results")
 
         if not results:
-            raise ValueError(f"location not found: {location}")
+            raise ToolError(f"location not found: {location}")
 
         place = results[0]
 
@@ -243,7 +322,7 @@ class WeatherTool(Tool):
             weather_codes.append(hourly["weather_code"][index])
 
         if not temperatures:
-            raise ValueError("no forecast data available for tomorrow")
+            raise ToolError("no forecast data available for tomorrow")
         return ToolResult(
             content=(
                 f"Location: {place['name']}, {place['country']}\n"
@@ -288,19 +367,19 @@ class FileSystemListTool(Tool):
         path = arguments.get("path")
 
         if not isinstance(path, str):
-            raise TypeError("path must be a string")
+            raise ToolError("path must be a string")
 
         sandbox = Path.home() / "Projects" / "personal-agent" / "data"
         target = (sandbox / path).resolve()
 
         if not target.is_relative_to(sandbox.resolve()):
-            raise ValueError("path is outside the filesystem sandbox")
+            raise ToolError("path is outside the filesystem sandbox")
 
         if not target.exists():
-            raise ValueError(f"path does not exist: {path}")
+            raise ToolError(f"path does not exist: {path}")
 
         if not target.is_dir():
-            raise ValueError(f"path is not a directory: {path}")
+            raise ToolError(f"path is not a directory: {path}")
 
         entries = []
 
@@ -342,28 +421,28 @@ class FileSystemReadTool(Tool):
         path = arguments.get("path")
 
         if not isinstance(path, str):
-            raise TypeError("path must be a string")
+            raise ToolError("path must be a string")
 
         sandbox = Path.home() / "Projects" / "personal-agent" / "data"
         sandbox = sandbox.resolve()
         target = (sandbox / path).resolve()
 
         if not target.is_relative_to(sandbox):
-            raise ValueError("path is outside the filesystem sandbox")
+            raise ToolError("path is outside the filesystem sandbox")
 
         if not target.exists():
-            raise ValueError(f"path does not exist: {path}")
+            raise ToolError(f"path does not exist: {path}")
 
         if not target.is_file():
-            raise ValueError(f"path is not a file: {path}")
+            raise ToolError(f"path is not a file: {path}")
 
         if target.stat().st_size > self.MAX_FILE_SIZE:
-            raise ValueError("file is too large to read")
+            raise ToolError("file is too large to read")
 
         try:
             content = target.read_text(encoding="utf-8")
-        except UnicodeDecodeError as error:
-            raise ValueError("file is not valid UTF-8 text") from error
+        except ToolError as error:
+            raise ToolError("file is not valid UTF-8 text") from error
 
         return ToolResult(content=content)