Added more memory management and fixed system prompt bug

8 files changed, 87 insertions, 9 deletions

diff --git a/personal_agent/agent.py b/personal_agent/agent.py
index 9c456a8..8efe949 100644
--- a/personal_agent/agent.py
+++ b/personal_agent/agent.py
@@ -59,6 +59,8 @@ class Agent:
             }
         )
 
+        # print(f"[debug: {self.context.build_messages(self.messages, self.working_memory)}]")
+
         for _ in range(5):
             response = self.llm.chat(
                 self.context.build_messages(
diff --git a/personal_agent/cli.py b/personal_agent/cli.py
index 57df7b0..a2d6af5 100644
--- a/personal_agent/cli.py
+++ b/personal_agent/cli.py
@@ -4,6 +4,7 @@ from personal_agent.agent import Agent
 from personal_agent.database.database import Database
 from personal_agent.identity import load_agent_identity, load_user_profile
 from personal_agent.llm import LLM
+from personal_agent.logging import configure_logging
 from personal_agent.memory.store import LongTermMemoryStore
 from personal_agent.tools import (
     CalculatorTool,
@@ -13,6 +14,8 @@ from personal_agent.tools import (
     WeatherTool,
 )
 
+configure_logging()
+
 llm = LLM(
     "http://localhost:8080/v1/chat/completions",
 )
diff --git a/personal_agent/context.py b/personal_agent/context.py
index 74cc254..defc835 100644
--- a/personal_agent/context.py
+++ b/personal_agent/context.py
@@ -1,6 +1,10 @@
+import logging
+
 from personal_agent.identity import AgentIdentity, UserProfile
 from personal_agent.memory.models import WorkingMemory
 
+logger = logging.getLogger(__name__)
+
 
 class ContextBuilder:
     def __init__(
@@ -18,21 +22,41 @@ class ContextBuilder:
         agent = self.agent_identity
         user = self.user_profile
 
-        return (
+        prompt = (
             f"You are {agent.name}, a personal AI assistant.\n"
             f"Personality: {agent.personality['style']}.\n"
             f"Honesty: {agent.values['honesty']}.\n"
             f"Directness: {agent.values['directness']}.\n"
             f"Favorite color: {agent.quirks['favorite_color']}.\n"
             f"\n"
+            f"===BEHAVIOR==="
             f"Task management:\n"
-            f"{agent.behavior['task_management']}"
+            f"{agent.behavior['task_management']}\n"
+            f"\n"
+            f"Working memory:\n"
+            f"{agent.behavior['working_memory']}"
+            f"\n"
+            f"Long term memory:\n"
+            f"{agent.behavior['long_term_memory']}"
+            f"===BEHAVIOR===\n"
             f"\n"
             f"The user's name is {user.name}"
             f"\n"
             f"{self.build_working_memory(memory)}"
         )
 
+        logger.debug("Built system prompt: %d characters", len(prompt))
+
+        logger.debug(
+            "Behavior section present: task_management=%s, working working_memory=%s, "
+            "long_term_memory=%s",
+            bool(agent.behavior.get("task_management")),
+            bool(agent.behavior.get("working_memory")),
+            bool(agent.behavior.get("long_term_memory")),
+        )
+
+        return prompt
+
     def build_working_memory(self, memory: WorkingMemory) -> str:
         lines = ["Current working memory:"]
 
diff --git a/personal_agent/database/migrations/002_add_memory_fts.sql b/personal_agent/database/migrations/002_add_memory_fts.sql
new file mode 100644
index 0000000..26ccada
--- /dev/null
+++ b/personal_agent/database/migrations/002_add_memory_fts.sql
@@ -0,0 +1,6 @@
+CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts
+USING fts5(
+    content,
+    content='memories',
+    content_rowid='rowid'
+);
diff --git a/personal_agent/database/migrations/003_populate_memory_fts.sql b/personal_agent/database/migrations/003_populate_memory_fts.sql
new file mode 100644
index 0000000..35c2a4e
--- /dev/null
+++ b/personal_agent/database/migrations/003_populate_memory_fts.sql
@@ -0,0 +1,3 @@
+INSERT INTO memories_fts(rowid, content)
+SELECT rowid, content
+FROM memories;
diff --git a/personal_agent/llm.py b/personal_agent/llm.py
index 418c478..cd35676 100644
--- a/personal_agent/llm.py
+++ b/personal_agent/llm.py
@@ -1,8 +1,11 @@
 import json
+import logging
 from dataclasses import dataclass
 
 import requests
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class ToolCall:
@@ -43,7 +46,8 @@ class LLM:
         response.raise_for_status()
 
         data = response.json()
-        print(f"\n[debug: {data}]\n")
+        logger.debug("LLM response: %s", data)
+        # print(f"\n[debug: {data}]\n")
         message = data["choices"][0]["message"]
 
         content = message.get("content")
diff --git a/personal_agent/logging.py b/personal_agent/logging.py
new file mode 100644
index 0000000..8f86003
--- /dev/null
+++ b/personal_agent/logging.py
@@ -0,0 +1,8 @@
+import logging
+
+
+def configure_logging(level: int = logging.INFO) -> None:
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    )
diff --git a/personal_agent/memory/store.py b/personal_agent/memory/store.py
index e5fca96..164d11f 100644
--- a/personal_agent/memory/store.py
+++ b/personal_agent/memory/store.py
@@ -21,6 +21,16 @@ class LongTermMemoryStore:
                 (str(memory.id), memory.content, memory.category.value, created_at),
             )
 
+            connection.execute(
+                """
+                INSERT INTO memories_fts(rowid, content)
+                SELECT rowid, content
+                FROM memories
+                WHERE id = ?
+                """,
+                (str(memory.id),),
+            )
+
     def get(self, memory_id: UUID) -> LongTermMemory | None:
         with self.database.connect() as connection:
             row = connection.execute(
@@ -52,6 +62,18 @@ class LongTermMemoryStore:
     def delete(self, memory_id: UUID) -> None:
         with self.database.connect() as connection:
             connection.execute(
+                """
+                DELETE FROM memories_fts
+                WHERE rowid = (
+                    SELECT rowid
+                    FROM memories
+                    WHERE id = ?
+                )
+                """,
+                (str(memory_id),),
+            )
+
+            connection.execute(
                 "DELETE FROM memories WHERE id = ?",
                 (str(memory_id),),
             )
@@ -68,18 +90,24 @@ class LongTermMemoryStore:
         if not query.strip():
             raise ValueError("Search query must not be empty")
 
+        terms = query.split()
+        fts_query = " OR ".join(
+            f'"{term.replace(chr(34), chr(34) * 2)}"' for term in terms
+        )
+
         sql = """
-            SELECT id, content, category, created_at
-            FROM memories
-            WHERE content LIKE ?
+            SELECT m.id, m.content, m.category, m.created_at
+            FROM memories_fts as f
+            JOIN memories AS m ON m.rowid = f.rowid
+            WHERE memories_fts MATCH ?
         """
-        parameters: list[str | int] = [f"%{query}%"]
+        parameters: list[str | int] = [fts_query]
 
         if category is not None:
-            sql += " AND category = ?"
+            sql += " AND m.category = ?"
             parameters.append(category.value)
 
-        sql += " ORDER BY id LIMIT ?"
+        sql += " ORDER BY bm25(memories_fts) LIMIT ?"
         parameters.append(limit)
 
         with self.database.connect() as connection: