From 97932dc44b06c3499d185e3f1cda162081bf2adf Mon Sep 17 00:00:00 2001
From: Yuhong Sun <yuhongsun96@gmail.com>
Date: Thu, 14 Nov 2024 17:28:03 -0800
Subject: [PATCH] Fix Quotes Prompting (#3137)

---
 backend/danswer/llm/answering/answer.py                    | 1 +
 backend/danswer/llm/answering/prompts/build.py             | 7 +++++++
 .../tools/tool_implementations/search_like_tool_utils.py   | 5 ++++-
 examples/assistants-api/topics_analyzer.py                 | 4 ++--
 4 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/backend/danswer/llm/answering/answer.py b/backend/danswer/llm/answering/answer.py
index c447db452..f9c9dbcb1 100644
--- a/backend/danswer/llm/answering/answer.py
+++ b/backend/danswer/llm/answering/answer.py
@@ -263,6 +263,7 @@ class Answer:
             message_history=self.message_history,
             llm_config=self.llm.config,
             single_message_history=self.single_message_history,
+            raw_user_text=self.question,
         )
         prompt_builder.update_system_prompt(
             default_build_system_message(self.prompt_config)
diff --git a/backend/danswer/llm/answering/prompts/build.py b/backend/danswer/llm/answering/prompts/build.py
index b5b774f52..ac9ce6f1a 100644
--- a/backend/danswer/llm/answering/prompts/build.py
+++ b/backend/danswer/llm/answering/prompts/build.py
@@ -59,6 +59,7 @@ class AnswerPromptBuilder:
         message_history: list[PreviousMessage],
         llm_config: LLMConfig,
         single_message_history: str | None = None,
+        raw_user_text: str | None = None,
     ) -> None:
         self.max_tokens = compute_max_llm_input_tokens(llm_config)
 
@@ -88,6 +89,12 @@ class AnswerPromptBuilder:
 
         self.new_messages_and_token_cnts: list[tuple[BaseMessage, int]] = []
 
+        self.raw_user_message = (
+            HumanMessage(content=raw_user_text)
+            if raw_user_text is not None
+            else user_message
+        )
+
     def update_system_prompt(self, system_message: SystemMessage | None) -> None:
         if not system_message:
             self.system_message_and_token_cnt = None
diff --git a/backend/danswer/tools/tool_implementations/search_like_tool_utils.py b/backend/danswer/tools/tool_implementations/search_like_tool_utils.py
index 6701f1602..121841d0b 100644
--- a/backend/danswer/tools/tool_implementations/search_like_tool_utils.py
+++ b/backend/danswer/tools/tool_implementations/search_like_tool_utils.py
@@ -55,9 +55,12 @@ def build_next_prompt_for_search_like_tool(
             )
         )
     elif answer_style_config.quotes_config:
+        # For Quotes, the system prompt is included in the user prompt
+        prompt_builder.update_system_prompt(None)
+
         prompt_builder.update_user_prompt(
             build_quotes_user_message(
-                message=prompt_builder.user_message_and_token_cnt[0],
+                message=prompt_builder.raw_user_message,
                 context_docs=final_context_documents,
                 history_str=prompt_builder.single_message_history or "",
                 prompt=prompt_config,
diff --git a/examples/assistants-api/topics_analyzer.py b/examples/assistants-api/topics_analyzer.py
index a8ef6c27a..d99c38436 100644
--- a/examples/assistants-api/topics_analyzer.py
+++ b/examples/assistants-api/topics_analyzer.py
@@ -28,7 +28,7 @@ If there is no relevant information, just say "No relevant information found."
 """
 
 
-def wait_on_run(client: OpenAI, run, thread):
+def wait_on_run(client: OpenAI, run, thread):  # type: ignore
     while run.status == "queued" or run.status == "in_progress":
         run = client.beta.threads.runs.retrieve(
             thread_id=thread.id,
@@ -38,7 +38,7 @@ def wait_on_run(client: OpenAI, run, thread):
     return run
 
 
-def show_response(messages) -> None:
+def show_response(messages) -> None:  # type: ignore
     # Get only the assistant's response text
     for message in messages.data[::-1]:
         if message.role == "assistant":