Merge pull request #123 from StampyAI/handle-discord

Handle Discord messages properly

api/src/stampy_chat/chat.py

@@ -12,7 +12,7 @@ from langchain.prompts import (
 from langchain.pydantic_v1 import Extra
 from langchain.schema import BaseMessage, ChatMessage, PromptValue, SystemMessage
 
-from stampy_chat.env import OPENAI_API_KEY
+from stampy_chat.env import OPENAI_API_KEY, COMPLETIONS_MODEL
 from stampy_chat.settings import Settings
 from stampy_chat.callbacks import StampyCallbackHandler, BroadcastCallbackHandler, LoggerCallbackHandler
 from stampy_chat.followups import StampyChain
@@ -142,7 +142,7 @@ def make_prompt(settings, chat_model, callbacks):
         example_prompt=ChatPromptTemplate.from_template(context_template, template_format="jinja2"),
         get_num_tokens=chat_model.get_num_tokens,
         max_tokens=settings.context_tokens,
-        input_variables=['query'],
+        input_variables=['query', 'history'],
     )
 
     # 2. The history items will be passed in from the memory
@@ -168,7 +168,7 @@ def make_prompt(settings, chat_model, callbacks):
 def make_memory(settings, history, callbacks):
     """Create a memory object to store the chat history."""
     memory = LimitedConversationSummaryBufferMemory(
-        llm=get_model(),
+        llm=get_model(model=COMPLETIONS_MODEL),  # used for summarization
         max_token_limit=settings.history_tokens,
         max_history=settings.maxHistory,
         chat_memory=ChatMessageHistory(),
@@ -179,6 +179,29 @@ def make_memory(settings, history, callbacks):
     return memory
 
 
+def merge_history(history):
+    """Merge subsequent messages into a single one.
+
+    ChatGPT works pretty much by alternating assistant and user queries. On the other
+    hand, systems like Slack or Discord will often have multiple messages as responses,
+    as people tend to write a few shorter messages rather than one big one. This function
+    will transform the later type of history into the former, so the LLM has an easier job.
+    """
+    if not history:
+        return history
+
+    messages = []
+    current_message = history[0]
+    for message in history[1:]:
+        if message.get('role') != current_message.get('role'):
+            messages.append(current_message)
+            current_message = message
+        else:
+            current_message['content'] += '\n' + message.get('content', '')
+    messages.append(current_message)
+    return messages
+
+
 def run_query(session_id: str, query: str, history: List[Dict], settings: Settings, callback: Callable[[Any], None] = None) -> Dict[str, str]:
     """Execute the query.
 
@@ -191,6 +214,8 @@ def run_query(session_id: str, query: str, history: List[Dict], settings: Settin
     callbacks = [LoggerCallbackHandler(session_id=session_id, query=query, history=history)]
     if callback:
         callbacks += [BroadcastCallbackHandler(callback)]
+
+    history = merge_history(history)
     chat_model = get_model(
         streaming=True,
         callbacks=callbacks,
@@ -205,7 +230,6 @@ def run_query(session_id: str, query: str, history: List[Dict], settings: Settin
         memory=make_memory(settings, history, callbacks)
     ) | StampyChain(callbacks=callbacks)
     result = chain.invoke({"query": query, 'history': history}, {'callbacks': []})
-
     if callback:
         callback({'state': 'done'})
         callback(None)  # make sure the callback handler know that things have ended

api/src/stampy_chat/citations.py

@@ -16,6 +16,8 @@ class ReferencesSelector(SemanticSimilarityExampleSelector):
     """Get examples with enumerated indexes added."""
 
     callbacks: List[StampyCallbackHandler] = []
+    history_field: str = 'history'
+    min_score: float = 0.8  # docs with lower scores will be excluded from the context
 
     class Config:
         """This is needed for extra fields to be added... """
@@ -27,6 +29,20 @@ class ReferencesSelector(SemanticSimilarityExampleSelector):
         """Make the reference used in citations - basically translate i -> 'a + i'"""
         return chr(i + 97)
 
+    def fetch_docs(self, input_variables) -> List:
+        ### Copied from parent - for some reason they ignore the ids of the returned items, so
+        # it has to be added manually here...
+        if self.input_keys:
+            input_variables = {key: input_variables[key] for key in self.input_keys}
+        query = " ".join(v for v in input_variables.values())
+        example_docs = [
+            doc for doc, score in self.vectorstore.similarity_search_with_score(query, k=self.k)
+            if score > self.min_score
+        ]
+
+        # Remove any duplicates - sometimes the same document is returned multiple times
+        return list({e.page_content: e for e in example_docs}.values())
+
     def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:
         """Fetch the top matching items from the underlying storage and add indexes.
 
@@ -36,22 +52,22 @@ class ReferencesSelector(SemanticSimilarityExampleSelector):
         for callback in self.callbacks:
             callback.on_context_fetch_start(input_variables)
 
-        ### Copied from parent - for some reason they ignore the ids of the returned items, so
-        # it has to be added manually here...
-        if self.input_keys:
-            input_variables = {key: input_variables[key] for key in self.input_keys}
-        query = " ".join(v for v in input_variables.values())
-        example_docs = self.vectorstore.similarity_search(query, k=self.k)
+        input_variables = dict(**input_variables)
+        history = input_variables.pop(self.history_field, [])
 
-        # Remove any duplicates - sometimes the same document is returned multiple times
-        example_docs = {e.page_content: e for e in example_docs}.values()
+        examples = self.fetch_docs(input_variables)
+
+        for item in history[::-1]:
+            if len(examples) >= self.k:
+                break
+            examples += self.fetch_docs({'answer': item.content})
 
         examples = [
             dict(
                 e.metadata,
                 id=e.page_content,
                 reference=self.make_reference(i)
-            ) for i, e in enumerate(example_docs)
+            ) for i, e in enumerate(examples)
         ]
 
         for callback in self.callbacks:
@@ -60,7 +76,7 @@ class ReferencesSelector(SemanticSimilarityExampleSelector):
         return examples
 
 
-def make_example_selector(k: int, **params) -> ReferencesSelector:
+def make_example_selector(**params) -> ReferencesSelector:
     embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
     vectorstore = Pinecone(PINECONE_INDEX, embeddings.embed_query, "hash_id", namespace=PINECONE_NAMESPACE)
     return ReferencesSelector(vectorstore=vectorstore, **params)

api/src/stampy_chat/settings.py

@@ -39,6 +39,10 @@ PROMPT_MODES = {
         "Put extra effort into explaining the intuition behind concepts "
         "rather than just giving a formal definition.\n\n"
     ),
+    "discord": (
+        "Your answer will be used in a Discord channel, so please Answer concisely, getting to "
+        "the crux of the matter in as few words as possible. Limit your answer to 1-2 paragraphs.\n\n"
+    ),
 }
 DEFAULT_PROMPTS = {
     'context': SOURCE_PROMPT,
@@ -148,7 +152,7 @@ class Settings:
 
     @property
     def mode_prompt(self):
-        return self.prompts['modes'].get(self.mode)
+        return self.prompts['modes'].get(self.mode, '')
 
     @property
     def question_prompt(self):

api/tests/stampy_chat/test_chat.py

@@ -11,6 +11,7 @@ from stampy_chat.chat import (
     MessageBufferPromptTemplate,
     PrefixedPrompt,
     make_memory,
+    merge_history,
 )
 
 
@@ -160,3 +161,43 @@ def test_make_memory_skips_deleted():
         ChatMessage(content='as should this', role='human'),
         ChatMessage(content='bla bla bla', role='assistant'),
     ])
+
+
+def test_merge_history_empty():
+    assert merge_history([]) == []
+
+
+def test_merge_history_no_merges():
+    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'},
+    ]
+    assert merge_history(history) == history
+
+
+def test_merge_history_merges():
+    history = [
+        {'role': 'user', 'content': 'question 1'},
+        {'role': 'assistant', 'content': 'answer 1 part 1'},
+        {'role': 'assistant', 'content': 'answer 1 part 2'},
+        {'role': 'user', 'content': 'question 2 part 1'},
+        {'role': 'user', 'content': 'question 2 part 2'},
+        {'role': 'user', 'content': 'question 2 part 3'},
+        {'role': 'assistant', 'content': 'answer 2'},
+        {'role': 'user', 'content': 'question 3'},
+        {'role': 'assistant', 'content': 'answer 3'},
+        {'role': 'user', 'content': 'question 4 part 1'},
+        {'role': 'user', 'content': 'question 4 part 2'},
+    ]
+    assert merge_history(history) == [
+        {'role': 'user', 'content': 'question 1'},
+        {'role': 'assistant', 'content': 'answer 1 part 1\nanswer 1 part 2'},
+        {'role': 'user', 'content': 'question 2 part 1\nquestion 2 part 2\nquestion 2 part 3'},
+        {'role': 'assistant', 'content': 'answer 2'},
+        {'role': 'user', 'content': 'question 3'},
+        {'role': 'assistant', 'content': 'answer 3'},
+        {'role': 'user', 'content': 'question 4 part 1\nquestion 4 part 2'},
+    ]

api/tests/stampy_chat/test_citations.py

@@ -20,6 +20,9 @@ class DummyVectorStore(VectorStore):
         pass
 
     def similarity_search(self, *args, **kwargs):
+        assert False, 'this should not have been called'
+
+    def similarity_search_with_score(self, *args, **kwargs):
         if self.similarity_search_return_value:
             return self.similarity_search_return_value
         elif self.similarity_search_func:
@@ -30,9 +33,9 @@ class DummyVectorStore(VectorStore):
 @pytest.fixture
 def selector():
     examples = [
-        Mock(page_content=f'{i}', metadata={
+        (Mock(page_content=f'{i}', metadata={
             'bla': f'bla {i}'
-        }) for i in range(5)
+        }), 0.81 + i / 10) for i in range(5)
     ]
     return ReferencesSelector(vectorstore=DummyVectorStore(similarity_search_return=examples))
 
@@ -75,9 +78,9 @@ def test_ReferencesSelector_select_examples_callbacks(selector):
 
 def test_ReferencesSelector_select_examples_removes_duplicates(selector):
     selector.vectorstore.similarity_search_return_value = [
-        Mock(page_content=f'{i}', metadata={
+        (Mock(page_content=f'{i}', metadata={
             'bla': f'bla {i}'
-        }) for i in range(5)
+        }), selector.min_score + 0.1 + i / 10) for i in range(5)
     ] * 5
 
     assert selector.select_examples(input_variables={}) == [
@@ -89,6 +92,73 @@ def test_ReferencesSelector_select_examples_removes_duplicates(selector):
     ]
 
 
+def test_ReferencesSelector_select_examples_removes_low_scores(selector):
+    def calc_score(i):
+        # odd numbers should have scores small enough to be removed
+        if i % 2 == 0:
+            return 0.81 + i / 10
+        else:
+            return 0.8 - i / 10
+
+    selector.vectorstore.similarity_search_return_value = [
+        (Mock(page_content=f'{i}', metadata={
+            'bla': f'bla {i}'
+        }), calc_score(i)) for i in range(5)
+    ]
+
+    assert selector.select_examples(input_variables={}) == [
+        {'bla': 'bla 0', 'id': '0', 'reference': 'a'},
+        {'bla': 'bla 2', 'id': '2', 'reference': 'b'},
+        {'bla': 'bla 4', 'id': '4', 'reference': 'c'},
+    ]
+
+
+def test_ReferencesSelector_select_examples_check_history(selector):
+
+    def searcher(query, *args, **kwargs):
+        return [(Mock(page_content=query, metadata={'bla': query}), 0.9)]
+
+    selector.vectorstore.similarity_search_return_value = None
+    selector.vectorstore.similarity_search_func = searcher
+
+    history = [
+        Mock(content='first history item'),
+        Mock(content='second history item'),
+        Mock(content='last history item'),
+    ]
+    assert selector.select_examples(input_variables={'query': 'queried value', 'history': history}) == [
+        {'bla': 'queried value', 'id': 'queried value', 'reference': 'a'},
+        {'bla': 'last history item', 'id': 'last history item', 'reference': 'b'},
+        {'bla': 'second history item', 'id': 'second history item', 'reference': 'c'},
+        {'bla': 'first history item', 'id': 'first history item', 'reference': 'd'}
+    ]
+
+
+def test_ReferencesSelector_select_examples_check_history_n_items(selector):
+    def searcher(query, *args, **kwargs):
+        return [
+            (Mock(page_content=f'{query} - {i}', metadata={'bla': query}), 0.9)
+            for i in range(3)
+        ]
+
+    selector.vectorstore.similarity_search_return_value = None
+    selector.vectorstore.similarity_search_func = searcher
+
+    history = [
+        Mock(content='first history item'),
+        Mock(content='second history item'),
+        Mock(content='last history item'),
+    ]
+    assert selector.select_examples(input_variables={'query': 'queried value', 'history': history}) == [
+        {'bla': 'queried value', 'id': 'queried value - 0', 'reference': 'a'},
+        {'bla': 'queried value', 'id': 'queried value - 1', 'reference': 'b'},
+        {'bla': 'queried value', 'id': 'queried value - 2', 'reference': 'c'},
+        {'bla': 'last history item', 'id': 'last history item - 0', 'reference': 'd'},
+        {'bla': 'last history item', 'id': 'last history item - 1', 'reference': 'e'},
+        {'bla': 'last history item', 'id': 'last history item - 2', 'reference': 'f'},
+    ]
+
+
 @pytest.mark.parametrize("overrides, expected", [
     # Basic fields
     ({}, {}),

web/src/hooks/useSettings.ts

@@ -27,6 +27,9 @@ const DEFAULT_PROMPTS = {
     'cite all of them. For example: "AGI is concerning [c, d, e]."\n\n',
   modes: {
     default: "",
+    discord:
+      "Your answer will be used in a Discord channel, so please Answer concisely, getting to " +
+      "the crux of the matter in as few words as possible. Limit your answer to 1-2 paragraphs.\n\n",
     concise:
       "Answer very concisely, getting to the crux of the matter in as " +
       "few words as possible. Limit your answer to 1-2 sentences.\n\n",
@@ -161,8 +164,12 @@ export default function useSettings() {
   };
 
   useEffect(() => {
-    const mode = (localStorage.getItem("chat_mode") as Mode) || "default";
-    updateSettings(makeSettings({ ...router.query, mode: mode }));
+    if (!router.isReady) return;
+
+    const mode = (router?.query?.mode ||
+      localStorage.getItem("chat_mode") ||
+      "default") as Mode;
+    updateSettings(makeSettings({ ...router.query, mode }));
     setLoaded(router.isReady);
   }, [router]);
 

web/src/types.ts

@@ -47,7 +47,7 @@ export type SearchResult = {
 };
 export type CurrentSearch = (AssistantEntry & { phase?: string }) | undefined;
 
-export type Mode = "rookie" | "concise" | "default";
+export type Mode = "rookie" | "concise" | "default" | "discord";
 
 export type LLMSettings = {
   prompts?: {