From 71879695b951dd30bbd6cbd25261e22ac7f8e5d1 Mon Sep 17 00:00:00 2001 From: batsumaru <> Date: Wed, 16 Sep 2026 17:33:42 +0900 Subject: Added more memory management and fixed system prompt bug --- tests/test_context.py | 128 +++++++++++++++++++++++++++++++++++++++++++++ tests/test_database.py | 8 +-- tests/test_memory_store.py | 59 +++++++++++++++++++++ 3 files changed, 191 insertions(+), 4 deletions(-) create mode 100644 tests/test_context.py (limited to 'tests') 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] -- cgit v1.3