Added more memory management and fixed system prompt bug

1 files changed, 128 insertions, 0 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