From 0a563fa8d6505f842811fec6be5c13708549ec35 Mon Sep 17 00:00:00 2001
From: Daniel O'Connell <github@ahiru.pl>
Date: Tue, 17 Oct 2023 20:05:20 +0200
Subject: [PATCH] Allow chat items to be deleted

---
 api/src/stampy_chat/chat.py        | 12 +----
 api/src/stampy_chat/db/session.py  | 16 +++----
 api/tests/stampy_chat/test_chat.py | 22 +++++++++-
 web/src/components/chat.tsx        | 23 ++++++++--
 web/src/components/entry.tsx       | 70 +++++++++++++-----------------
 web/src/hooks/useSearch.ts         |  5 ++-
 web/src/types.ts                   |  4 ++
 7 files changed, 85 insertions(+), 67 deletions(-)

diff --git a/api/src/stampy_chat/chat.py b/api/src/stampy_chat/chat.py
index d70159e..a88c127 100644
--- a/api/src/stampy_chat/chat.py
+++ b/api/src/stampy_chat/chat.py
@@ -175,7 +175,7 @@ def make_memory(settings, history, callbacks):
         return_messages=True,
         callbacks=callbacks
     )
-    memory.set_messages(history)
+    memory.set_messages([i for i in history if i.get('role') != 'deleted'])
     return memory
 
 
@@ -193,16 +193,6 @@ def run_query(session_id: str, query: str, history: List[Dict], settings: Settin
         callbacks += [BroadcastCallbackHandler(callback)]
     chat_model = get_model(streaming=True, callbacks=callbacks, max_tokens=settings.max_response_tokens)
 
-    memory = LimitedConversationSummaryBufferMemory(
-        llm=get_model(),
-        max_token_limit=settings.history_tokens,
-        max_history=settings.maxHistory,
-        chat_memory=ChatMessageHistory(),
-        return_messages=True,
-        callbacks=callbacks
-    )
-    memory.set_messages(history)
-
     chain = LLMChain(
         llm=chat_model,
         verbose=False,
diff --git a/api/src/stampy_chat/db/session.py b/api/src/stampy_chat/db/session.py
index 4d99173..6342d70 100644
--- a/api/src/stampy_chat/db/session.py
+++ b/api/src/stampy_chat/db/session.py
@@ -48,17 +48,15 @@ class ItemAdder:
         self._last_save = time.time()
 
     def commit(self):
-        with Session(self.engine) as session:
-            try:
+        try:
+            with Session(self.engine) as session:
                 session.add_all(self.batch)
                 session.commit()
                 logger.debug('added %s items', len(self.batch))
-                self.batch = []
-            except SQLAlchemyError as e:
-                logger.warn('Got error when trying to commit to database: %s', e)
-                session.rollback()
-                raise e
+            self.batch = []
             self._last_save = time.time()
+        except SQLAlchemyError as e:
+            logger.warn('Got error when trying to commit to database: %s', e)
 
     def add(self, *items):
         """Add the provided items to the database, commiting them if needed."""
@@ -69,6 +67,4 @@ class ItemAdder:
 
     def __del__(self):
         logger.debug('cleaning up session')
-        if self.session:
-            self.commit()
-            self.session.close()
+        self.commit()
diff --git a/api/tests/stampy_chat/test_chat.py b/api/tests/stampy_chat/test_chat.py
index dfb51bf..368ac2a 100644
--- a/api/tests/stampy_chat/test_chat.py
+++ b/api/tests/stampy_chat/test_chat.py
@@ -1,13 +1,16 @@
+from unittest.mock import patch
 from langchain.llms.fake import FakeListLLM
 from langchain.memory import ChatMessageHistory
 from langchain.prompts import ChatPromptTemplate
 from langchain.schema import ChatMessage, HumanMessage, SystemMessage
 
+from stampy_chat.settings import Settings
 from stampy_chat.callbacks import StampyCallbackHandler
 from stampy_chat.chat import (
     LimitedConversationSummaryBufferMemory,
     MessageBufferPromptTemplate,
-    PrefixedPrompt
+    PrefixedPrompt,
+    make_memory,
 )
 
 
@@ -140,3 +143,20 @@ def test_LimitedConversationSummaryBufferMemory_set_with_callbacks():
         'start': history,
         'end': memory.chat_memory,
     }
+
+
+def test_make_memory_skips_deleted():
+    history = [
+        {'content': 'this should be kept', 'role': 'system'},
+        {'content': 'as should this', 'role': 'human'},
+        {'content': 'this will be ignored', 'role': 'deleted'},
+        {'content': 'bla bla bla', 'role': 'assistant'},
+        {'content': 'remove me!!', 'role': 'deleted'},
+    ]
+    with patch('stampy_chat.chat.get_model', return_value=FakeListLLM(responses=[])):
+        mem = make_memory(Settings(), history, [])
+    assert mem.chat_memory == ChatMessageHistory(messages=[
+        ChatMessage(content='this should be kept', role='system'),
+        ChatMessage(content='as should this', role='human'),
+        ChatMessage(content='bla bla bla', role='assistant'),
+    ])
diff --git a/web/src/components/chat.tsx b/web/src/components/chat.tsx
index 1987728..bedaa92 100644
--- a/web/src/components/chat.tsx
+++ b/web/src/components/chat.tsx
@@ -164,9 +164,26 @@ const Chat = ({ sessionId, settings, onQuery, onNewEntry }: ChatParams) => {
 
   return (
     <ul className="flex-auto">
-      {entries.map((entry, i) => (
-        <EntryTag entry={entry} key={i} />
-      ))}
+      {entries.map(
+        (entry, i) =>
+          !entry.deleted && (
+            <li className="group relative flex" key={i}>
+              <EntryTag entry={entry} />
+              <span
+                className="delete-item absolute right-5 hidden cursor-pointer group-hover:block"
+                onClick={() => {
+                  const entry = entries[i];
+                  if (entry !== undefined) {
+                    entry.deleted = true;
+                    setEntries([...entries]);
+                  }
+                }}
+              >
+                X
+              </span>
+            </li>
+          )
+      )}
       <SearchBox search={search} onQuery={onQuery} />
 
       {last_entry}
diff --git a/web/src/components/entry.tsx b/web/src/components/entry.tsx
index 2e052b6..5dc28dd 100644
--- a/web/src/components/entry.tsx
+++ b/web/src/components/entry.tsx
@@ -13,60 +13,50 @@ import TextareaAutosize from "react-textarea-autosize";
 
 export const User = ({ entry }: { entry: UserEntry }) => {
   return (
-    <li className="mt-1 mb-2 flex">
-      <TextareaAutosize
-        className="flex-1 resize-none border border-gray-300 px-1"
-        value={entry.content}
-      />
-    </li>
+    <TextareaAutosize
+      className="flex-1 resize-none border border-gray-300 px-1"
+      value={entry.content}
+    />
   );
 };
 
 export const Error = ({ entry }: { entry: ErrorMessage }) => {
   return (
-    <li>
-      <p className="border border-red-500 bg-red-100 px-1 text-red-800">
-        {" "}
-        {entry.content}{" "}
-      </p>
-    </li>
+    <p className="border border-red-500 bg-red-100 px-1 text-red-800">
+      {" "}
+      {entry.content}{" "}
+    </p>
   );
 };
 
 export const Assistant = ({ entry }: { entry: AssistantEntryType }) => {
-  return (
-    <li>
-      <AssistantEntry entry={entry} />
-    </li>
-  );
+  return <AssistantEntry entry={entry} />;
 };
 
 export const Stampy = ({ entry }: { entry: StampyMessage }) => {
   return (
-    <li>
-      <div
-        className="my-7 rounded bg-slate-500 px-4 py-0.5 text-slate-50"
-        style={{
-          marginLeft: "auto",
-          marginRight: "auto",
-          maxWidth: "99.8%",
-        }}
-      >
-        <div>
-          <GlossarySpan content={entry.content} />
-        </div>
-        <div className="mb-3 flex justify-end">
-          <a
-            href={entry.url}
-            target="_blank"
-            className="flex items-center space-x-1"
-          >
-            <span>aisafety.info</span>
-            <Image src={logo} alt="aisafety.info logo" width={19} />
-          </a>
-        </div>
+    <div
+      className="my-7 rounded bg-slate-500 px-4 py-0.5 text-slate-50"
+      style={{
+        marginLeft: "auto",
+        marginRight: "auto",
+        maxWidth: "99.8%",
+      }}
+    >
+      <div>
+        <GlossarySpan content={entry.content} />
       </div>
-    </li>
+      <div className="mb-3 flex justify-end">
+        <a
+          href={entry.url}
+          target="_blank"
+          className="flex items-center space-x-1"
+        >
+          <span>aisafety.info</span>
+          <Image src={logo} alt="aisafety.info logo" width={19} />
+        </a>
+      </div>
+    </div>
   );
 };
 
diff --git a/web/src/hooks/useSearch.ts b/web/src/hooks/useSearch.ts
index cfc7777..de0a62a 100644
--- a/web/src/hooks/useSearch.ts
+++ b/web/src/hooks/useSearch.ts
@@ -16,8 +16,9 @@ const MAX_FOLLOWUPS = 4;
 const DATA_HEADER = "data: ";
 const EVENT_END_HEADER = "event: close";
 
+type EntryRole = "error" | "stampy" | "assistant" | "user" | "deleted";
 type HistoryEntry = {
-  role: "error" | "stampy" | "assistant" | "user";
+  role: EntryRole;
   content: string;
 };
 
@@ -223,7 +224,7 @@ export const runSearch = async (
     const history = entries
       .filter((entry) => entry.role !== "error")
       .map((entry) => ({
-        role: entry.role,
+        role: (entry.deleted ? "deleted" : entry.role) as EntryRole,
         content: entry.content.trim(),
       }));
 
diff --git a/web/src/types.ts b/web/src/types.ts
index b5d8b45..aa99fc1 100644
--- a/web/src/types.ts
+++ b/web/src/types.ts
@@ -17,6 +17,7 @@ export type Entry = UserEntry | AssistantEntry | ErrorMessage | StampyMessage;
 export type UserEntry = {
   role: "user";
   content: string;
+  deleted?: boolean;
 };
 
 export type AssistantEntry = {
@@ -24,17 +25,20 @@ export type AssistantEntry = {
   content: string;
   citations: Citation[];
   citationsMap: Map<string, Citation>;
+  deleted?: boolean;
 };
 
 export type ErrorMessage = {
   role: "error";
   content: string;
+  deleted?: boolean;
 };
 
 export type StampyMessage = {
   role: "stampy";
   content: string;
   url: string;
+  deleted?: boolean;
 };
 
 export type SearchResult = {