Merge pull request #6 from crux-bphc/memory-fix

Memory Bug Fix
crux-bphc · Jan 18, 2025 · 7a99c6f · 7a99c6f
2 parents 999502f + 35fb29d
commit 7a99c6f
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 10 deletions.
diff --git a/src/memory/long_term_memory.py b/src/memory/long_term_memory.py
@@ -38,7 +38,7 @@ def add_long_term_memory(
     cur = conn.cursor()
     cur.execute(
         """
-        CREATE TABLE IF NOT EXISTS longterm_memory (
+        CREATE TABLE IF NOT EXISTS long_term_memory (
             id VARCHAR(255) NOT NULL,
             memory VARCHAR(255) NOT NULL,
             category VARCHAR(255) NOT NULL
@@ -48,7 +48,7 @@ def add_long_term_memory(
     if action == "Create":
         cur.execute(
             f"""
-            INSERT INTO longterm_memory (id, memory, category)
+            INSERT INTO long_term_memory (id, memory, category)
             VALUES ('{id}', '{memory}', '{Category(category).value}');
             """
         )
@@ -57,7 +57,7 @@ def add_long_term_memory(
     elif action == "Update":
         cur.execute(
             f"""
-            UPDATE longterm_memory
+            UPDATE long_term_memory
             SET memory = '{memory}'
             WHERE id = '{id}' AND memory = '{memory_old}';
             """
@@ -67,7 +67,7 @@ def add_long_term_memory(
     elif action == "Delete":
         cur.execute(
             f"""
-            DELETE FROM longterm_memory
+            DELETE FROM long_term_memory
             WHERE id = '{id}' AND memory = '{memory_old}' AND category = '{Category(category).value}';
             """
         )
@@ -85,7 +85,7 @@ def fetch_long_term_memory(id: str) -> list[LongTermMemory]:
     cur.execute(
         f"""
         SELECT memory, category
-        FROM longterm_memory
+        FROM long_term_memory
         WHERE id = '{id}';
         """
     )
@@ -104,7 +104,7 @@ def reset_long_term_memory(id: str) -> None:
     cur = conn.cursor()
     cur.execute(
         f"""
-        DELETE FROM longterm_memory
+        DELETE FROM long_term_memory
         WHERE id = '{id}';
         """
     )

diff --git a/src/memory/short_term_memory.py b/src/memory/short_term_memory.py
@@ -26,11 +26,15 @@ def add_short_term_memory(id: str, query: str, reply: str, agent: str) -> None:
     cur = conn.cursor()
     cur.execute(
         """
-        CREATE TABLE IF NOT EXISTS short_term_memory (id VARCHAR(255) NOT NULL, message_id VARCHAR(255) NOT NULL, created_at TIMESTAMP NOT NULL, query VARCHAR(255) NOT NULL, reply VARCHAR(255) NOT NULL, agent VARCHAR(255) NOT NULL);
+        CREATE TABLE IF NOT EXISTS short_term_memory (id VARCHAR(255) NOT NULL, message_id VARCHAR(255) NOT NULL, created_at TIMESTAMP NOT NULL, query TEXT NOT NULL, reply TEXT NOT NULL, agent VARCHAR(255) NOT NULL);
         """
     )
     message_id = str(uuid.uuid4())
     created_at = datetime.now()
+
+    query = query.replace("'", "''")
+    reply = reply.replace("'", "''")
+
     cur.execute(
         f"""
         INSERT INTO short_term_memory (id, message_id, created_at, query, reply, agent)

diff --git a/src/nodes.py b/src/nodes.py
@@ -25,9 +25,7 @@ def course_query(state: State):
 def general_campus_query(state: State):
     query = state["messages"][0].content
     result = agents.general_campus_query(query, state.get("chat_history", ""))
-    add_short_term_memory(
-        state["user_id"], query, result.content, "general_campus_query"
-    )
+    add_short_term_memory(state["user_id"], query, result, "general_campus_query")
     return {"messages": [result]}