diff --git a/api/src/stampy_chat/chat.py b/api/src/stampy_chat/chat.py index 5d7f8f4..a1f3786 100644 --- a/api/src/stampy_chat/chat.py +++ b/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 diff --git a/api/src/stampy_chat/citations.py b/api/src/stampy_chat/citations.py index 3b8d706..8c9be85 100644 --- a/api/src/stampy_chat/citations.py +++ b/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) diff --git a/api/src/stampy_chat/settings.py b/api/src/stampy_chat/settings.py index 00af508..12aae79 100644 --- a/api/src/stampy_chat/settings.py +++ b/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): diff --git a/api/tests/stampy_chat/test_chat.py b/api/tests/stampy_chat/test_chat.py index 368ac2a..29c0383 100644 --- a/api/tests/stampy_chat/test_chat.py +++ b/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'}, + ] diff --git a/api/tests/stampy_chat/test_citations.py b/api/tests/stampy_chat/test_citations.py index 024bad9..a9e8925 100644 --- a/api/tests/stampy_chat/test_citations.py +++ b/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 ({}, {}), diff --git a/web/src/hooks/useSettings.ts b/web/src/hooks/useSettings.ts index c0f642d..09efb51 100644 --- a/web/src/hooks/useSettings.ts +++ b/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]); diff --git a/web/src/types.ts b/web/src/types.ts index e751c66..537d40c 100644 --- a/web/src/types.ts +++ b/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?: {