Added more memory management and fixed system prompt bug

3 files changed, 191 insertions, 4 deletions

diff --git a/tests/test_context.py b/tests/test_context.py
new file mode 100644
index 0000000..baf5dd6
--- /dev/null
+++ b/tests/test_context.py
@@ -0,0 +1,128 @@
+from personal_agent.context import ContextBuilder
+from personal_agent.identity import AgentIdentity, UserProfile
+from personal_agent.memory.models import TaskStatus, WorkingMemory
+
+
+def create_agent_identity() -> AgentIdentity:
+    return AgentIdentity(
+        name="Test Agent",
+        personality={"style": "direct"},
+        values={"honesty": True, "directness": True},
+        quirks={"favorite_color": "blue"},
+        behavior={
+            "task_management": "TASK MANAGEMENT POLICY",
+            "working_memory": "WORKING MEMORY POLICY",
+            "long_term_memory": "LONG TERM MEMORY POLICY",
+        },
+    )
+
+
+def create_user_profile() -> UserProfile:
+    return UserProfile(
+        name="Test User",
+        preferences={},
+    )
+
+
+def create_context_builder() -> ContextBuilder:
+    return ContextBuilder(
+        agent_identity=create_agent_identity(),
+        user_profile=create_user_profile(),
+    )
+
+
+def test_system_prompt_contains_agent_identity():
+    builder = create_context_builder()
+
+    prompt = builder.build_system_prompt(WorkingMemory())
+
+    assert "You are Test Agent, a personal AI assistant." in prompt
+    assert "Personality: direct." in prompt
+    assert "Honesty: True." in prompt
+    assert "Directness: True." in prompt
+    assert "Favorite color: blue." in prompt
+    assert "The user's name is Test User" in prompt
+
+
+def test_system_prompt_contains_all_behavior_policies():
+    builder = create_context_builder()
+
+    prompt = builder.build_system_prompt(WorkingMemory())
+
+    behavior = builder.agent_identity.behavior
+
+    assert behavior["task_management"] in prompt
+    assert behavior["working_memory"] in prompt
+    assert behavior["long_term_memory"] in prompt
+
+
+def test_system_prompt_contains_behavior_sections():
+    builder = create_context_builder()
+
+    prompt = builder.build_system_prompt(WorkingMemory())
+
+    assert "===BEHAVIOR===" in prompt
+    assert "Task management:" in prompt
+    assert "Working memory:" in prompt
+    assert "Long term memory:" in prompt
+
+
+def test_system_prompt_contains_working_memory():
+    builder = create_context_builder()
+
+    memory = WorkingMemory()
+    memory.set_objective("Implement authentication")
+    memory.set_status(TaskStatus.ACTIVE)
+    memory.add_fact("The project uses SQLite.")
+    memory.add_note("Need to add logout support.")
+
+    prompt = builder.build_system_prompt(memory)
+
+    assert "Current working memory:" in prompt
+    assert "Objective: Implement authentication" in prompt
+    assert "Status: active" in prompt
+    assert "The project uses SQLite." in prompt
+    assert "Need to add logout support." in prompt
+
+
+def test_empty_working_memory_is_explicitly_shown():
+    builder = create_context_builder()
+
+    prompt = builder.build_system_prompt(WorkingMemory())
+
+    assert "Current working memory:" in prompt
+    assert "(empty)" in prompt
+
+
+def test_build_messages_puts_system_prompt_first():
+    builder = create_context_builder()
+
+    history: list[dict[str, object]] = [
+        {
+            "role": "user",
+            "content": "Hello",
+        },
+        {
+            "role": "assistant",
+            "content": "Hi!",
+        },
+    ]
+
+    messages = builder.build_messages(history, WorkingMemory())
+
+    assert messages[0]["role"] == "system"
+    assert messages[1:] == history
+
+
+def test_build_messages_uses_current_working_memory():
+    builder = create_context_builder()
+
+    memory = WorkingMemory()
+    memory.set_objective("Build the memory system")
+
+    messages = builder.build_messages([], memory)
+
+    system_prompt = messages[0]["content"]
+
+    assert isinstance(system_prompt, str)
+    assert "Objective: Build the memory system" in system_prompt
diff --git a/tests/test_database.py b/tests/test_database.py
index 290b985..e387123 100644
--- a/tests/test_database.py
+++ b/tests/test_database.py
@@ -34,7 +34,7 @@ def test_migrate_records_version(tmp_path):
     with database.connect() as connection:
         rows = connection.execute("SELECT version FROM schema_migrations").fetchall()
 
-    assert rows == [(1,)]
+    assert rows == [(1,), (2,), (3,)]
 
 
 def test_migrate_is_idempotent(tmp_path):
@@ -46,13 +46,13 @@ def test_migrate_is_idempotent(tmp_path):
     with database.connect() as connection:
         rows = connection.execute("SELECT version FROM schema_migrations").fetchall()
 
-    assert rows == [(1,)]
+    assert rows == [(1,), (2,), (3,)]
 
 
 def test_failed_migration_rolls_back(tmp_path):
     database = Database(tmp_path / "test.db")
 
-    migration = database.migrations_dir / "002_failing_migration.sql"
+    migration = database.migrations_dir / "999_failing_migration.sql"
     migration.write_text(
         """
         THIS IS INVALID SQL;
@@ -68,7 +68,7 @@ def test_failed_migration_rolls_back(tmp_path):
                 "SELECT version FROM schema_migrations"
             ).fetchall()
 
-        assert migrations == [(1,)]
+        assert migrations == [(1,), (2,), (3,)]
     finally:
         migration.unlink()
 
diff --git a/tests/test_memory_store.py b/tests/test_memory_store.py
index 8f3407d..b83d217 100644
--- a/tests/test_memory_store.py
+++ b/tests/test_memory_store.py
@@ -211,3 +211,62 @@ def test_search_respects_limit(tmp_path):
         store.add(memory)
 
     assert store.search("programming", limit=2) == memories[:2]
+
+
+def test_search_uses_fts(tmp_path):
+    database = Database(tmp_path / "agent.db")
+    database.migrate()
+    store = LongTermMemoryStore(database)
+
+    memory = LongTermMemory.create(
+        "User prefers Rust for systems programming.",
+        MemoryCategory.SEMANTIC,
+    )
+    store.add(memory)
+
+    results = store.search("preferred systems programming language")
+
+    assert [result.id for result in results] == [memory.id]
+
+
+def test_delete_removes_memory_from_search(tmp_path):
+    database = Database(tmp_path / "agent.db")
+    database.migrate()
+    store = LongTermMemoryStore(database)
+
+    memory = LongTermMemory.create(
+        "User prefers Rust for systems programming.",
+        MemoryCategory.SEMANTIC,
+    )
+    store.add(memory)
+
+    assert store.search("Rust")
+
+    store.delete(memory.id)
+
+    assert store.search("Rust") == []
+
+
+def test_search_filters_by_category(tmp_path):
+    database = Database(tmp_path / "agent.db")
+    database.migrate()
+    store = LongTermMemoryStore(database)
+
+    semantic = LongTermMemory.create(
+        "User prefers Rust.",
+        MemoryCategory.SEMANTIC,
+    )
+    episodic = LongTermMemory.create(
+        "User used Rust on a filesystem project.",
+        MemoryCategory.EPISODIC,
+    )
+
+    store.add(semantic)
+    store.add(episodic)
+
+    results = store.search(
+        "Rust",
+        category=MemoryCategory.EPISODIC,
+    )
+
+    assert [result.id for result in results] == [episodic.id]