proper order of messages

Author: joachim-danswer

backend/onyx/agents/agent_search/dr/nodes/dr_a0_clarification.py

@@ -84,6 +84,10 @@
 
 logger = setup_logger()
 
+_ANSWER_COMMENT_PROMPT = "I will now answer your question directly."
+
+_CONSIDER_TOOLS_PROMPT = "I will now concier the tools and sub-agents that are available to answer your question."
+
 
 def _format_tool_name(tool_name: str) -> str:
     """Convert tool name to LLM-friendly format."""
@@ -432,25 +436,19 @@ def clarifier(
         assistant_system_prompt = PromptTemplate(DEFAULT_DR_SYSTEM_PROMPT).build()
         assistant_task_prompt = ""
 
-    chat_history_string = (
-        get_chat_history_string(
-            graph_config.inputs.prompt_builder.message_history,
-            MAX_CHAT_HISTORY_MESSAGES,
-        )
-        or "(No chat history yet available)"
-    )
+    # chat_history_string = (
+    #     get_chat_history_string(
+    #         graph_config.inputs.prompt_builder.message_history,
+    #         MAX_CHAT_HISTORY_MESSAGES,
+    #     )
+    #     or "(No chat history yet available)"
+    # )
 
     chat_history_messages = get_chat_history_messages(
-        graph_config.inputs.prompt_builder.message_history, MAX_CHAT_HISTORY_MESSAGES
+        graph_config.inputs.prompt_builder.raw_message_history,
+        MAX_CHAT_HISTORY_MESSAGES,
     )
 
-    if len(chat_history_messages) > 0:
-        chat_history_messages = [
-            SystemMessage(content="Here are the previous messages in the chat history:")
-        ] + chat_history_messages
-    else:
-        chat_history_messages = []
-
     uploaded_text_context = (
         _construct_uploaded_text_context(graph_config.inputs.files)
         if graph_config.inputs.files
@@ -495,7 +493,21 @@ def clarifier(
     message_history_for_continuation.append(SystemMessage(content=base_system_message))
     message_history_for_continuation.extend(chat_history_messages)
     message_history_for_continuation.extend(uploaded_file_messages)
-    message_history_for_continuation.append(HumanMessage(content=original_question))
+
+    # Create message content that includes text and any available images
+    message_content: list[dict[str, Any]] = [
+        {"type": "text", "text": original_question}
+    ]
+    if uploaded_image_context:
+        message_content.extend(uploaded_image_context)
+
+    # If we only have text, use string content for backwards compatibility
+    if len(message_content) == 1:
+        message_history_for_continuation.append(HumanMessage(content=original_question))
+    else:
+        message_history_for_continuation.append(
+            HumanMessage(content=cast(list[str | dict[Any, Any]], message_content))
+        )
     message_history_for_continuation.append(AIMessage(content=QUESTION_CONFIRMATION))
 
     if not (force_use_tool and force_use_tool.force_use):
@@ -508,7 +520,7 @@ def clarifier(
         if not use_tool_calling_llm or len(available_tools) == 1:
             if len(available_tools) > 1:
                 message_history_for_continuation.append(
-                    SystemMessage(content=DECISION_PROMPT_WO_TOOL_CALLING)
+                    HumanMessage(content=DECISION_PROMPT_WO_TOOL_CALLING)
                 )
 
                 llm_decision = invoke_llm_json(
@@ -535,6 +547,10 @@ def clarifier(
                     reminder=reminder,
                 )
 
+                message_history_for_continuation.append(
+                    AIMessage(content=_ANSWER_COMMENT_PROMPT)
+                )
+
                 message_history_for_continuation.append(
                     HumanMessage(content=answer_prompt)
                 )
@@ -795,14 +811,7 @@ def clarifier(
     else:
         next_tool = DRPath.ORCHESTRATOR.value
 
-    if research_type == ResearchType.DEEP and clarification:
-        message_history_for_continuation.append(
-            AIMessage(content=clarification.clarification_question)
-        )
-        if clarification.clarification_response:
-            message_history_for_continuation.append(
-                HumanMessage(content=clarification.clarification_response)
-            )
+    message_history_for_continuation.append(AIMessage(content=_CONSIDER_TOOLS_PROMPT))
 
     return OrchestrationSetup(
         original_question=original_question,

backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py

@@ -36,7 +36,6 @@
 from onyx.configs.agent_configs import TF_DR_TIMEOUT_LONG
 from onyx.configs.agent_configs import TF_DR_TIMEOUT_SHORT
 from onyx.prompts.dr_prompts import DEFAULLT_DECISION_PROMPT
-from onyx.prompts.dr_prompts import NEXT_TOOL_PURPOSE_PROMPT
 from onyx.prompts.dr_prompts import REPEAT_PROMPT
 from onyx.prompts.dr_prompts import SUFFICIENT_INFORMATION_STRING
 from onyx.prompts.dr_prompts import TOOL_CHOICE_WRAPPER_PROMPT
@@ -50,6 +49,10 @@
 _DECISION_SYSTEM_PROMPT_PREFIX = "Here are general instructions by the user, which \
 may or may not influence the decision what to do next:\n\n"
 
+_PLAN_OF_RECORD_PROMPT = "Can you create a plan of record?"
+
+_NEXT_ACTION_PROMPT = "What should be the next action?"
+
 
 def _get_implied_next_tool_based_on_tool_call_history(
     tools_used: list[str],
@@ -86,7 +89,7 @@ def orchestrator(
     clarification = state.clarification
     assistant_system_prompt = state.assistant_system_prompt
 
-    message_history_for_continuation = state.orchestration_llm_messages
+    message_history_for_continuation = list(state.orchestration_llm_messages)
     new_messages: list[SystemMessage | HumanMessage | AIMessage] = []
 
     if assistant_system_prompt:
@@ -126,24 +129,19 @@ def orchestrator(
         or "(No answer history yet available)"
     )
 
-    if (
-        research_type == ResearchType.DEEP
-        and most_recent_answer_history_wo_docs_string
-        != "(No answer history yet available)"
-    ):
-        message_history_for_continuation.append(
-            AIMessage(content=most_recent_answer_history_wo_docs_string)
-        )
-        new_messages.append(
-            AIMessage(content=most_recent_answer_history_wo_docs_string)
-        )
-    elif (
-        most_recent_answer_history_w_docs_string != "(No answer history yet available)"
-    ):
-        message_history_for_continuation.append(
-            AIMessage(content=most_recent_answer_history_w_docs_string)
-        )
-        new_messages.append(AIMessage(content=most_recent_answer_history_w_docs_string))
+    human_text = ai_text = ""
+    if most_recent_answer_history_wo_docs_string != "(No answer history yet available)":
+        human_text = f"Results from Iteration {iteration_nr - 1}?"
+        if research_type == ResearchType.DEEP:
+            ai_text = most_recent_answer_history_wo_docs_string
+        else:
+            ai_text = most_recent_answer_history_w_docs_string
+
+        message_history_for_continuation.append(HumanMessage(content=human_text))
+        new_messages.append(HumanMessage(content=human_text))
+
+        message_history_for_continuation.append(AIMessage(content=ai_text))
+        new_messages.append(AIMessage(content=ai_text))
 
     next_tool_name = None
 
@@ -367,8 +365,6 @@ def orchestrator(
             uploaded_context=uploaded_context,
         )
 
-        message_history_for_continuation.append(HumanMessage(content=decision_prompt))
-
         if remaining_time_budget > 0:
             try:
                 orchestrator_action = invoke_llm_json(
@@ -440,6 +436,18 @@ def orchestrator(
                 writer,
             )
 
+            message_history_for_continuation.append(
+                HumanMessage(content=_PLAN_OF_RECORD_PROMPT)
+            )
+            new_messages.append(HumanMessage(content=_PLAN_OF_RECORD_PROMPT))
+
+            message_history_for_continuation.append(
+                AIMessage(content=f"{HIGH_LEVEL_PLAN_PREFIX}\n\n {plan_of_record.plan}")
+            )
+            new_messages.append(
+                AIMessage(content=f"{HIGH_LEVEL_PLAN_PREFIX}\n\n {plan_of_record.plan}")
+            )
+
             start_time = datetime.now()
 
             repeat_plan_prompt = REPEAT_PROMPT.build(
@@ -573,16 +581,14 @@ def orchestrator(
         questions="\n - " + "\n - ".join(query_list or []),
     )
 
+    message_history_for_continuation.append(HumanMessage(content=_NEXT_ACTION_PROMPT))
+    new_messages.append(HumanMessage(content=_NEXT_ACTION_PROMPT))
+
     message_history_for_continuation.append(
         AIMessage(content=tool_choice_wrapper_prompt)
     )
     new_messages.append(AIMessage(content=tool_choice_wrapper_prompt))
 
-    message_history_for_continuation.append(
-        HumanMessage(content=NEXT_TOOL_PURPOSE_PROMPT)
-    )
-    new_messages.append(HumanMessage(content=NEXT_TOOL_PURPOSE_PROMPT))
-
     purpose_tokens: list[str] = [""]
     purpose = ""
 

backend/onyx/agents/agent_search/dr/nodes/dr_a2_closer.py

@@ -55,8 +55,11 @@
 from onyx.utils.logger import setup_logger
 from onyx.utils.threadpool_concurrency import run_with_timeout
 
+
 logger = setup_logger()
 
+_SOURCE_MATERIAL_PROMPT = "Can yut please put together all of the supporting material?"
+
 
 def extract_citation_numbers(text: str) -> list[int]:
     """
@@ -334,6 +337,10 @@ def closer(
             final_user_instructions=assistant_task_prompt
             or "(No final user instructions)",
         )
+
+        message_history_for_final_answer.append(
+            HumanMessage(content=_SOURCE_MATERIAL_PROMPT)
+        )
         message_history_for_final_answer.append(
             AIMessage(
                 content=FINAL_ANSWER_DEEP_CITATION_PROMPT.build(

backend/onyx/agents/agent_search/dr/utils.py

@@ -4,7 +4,6 @@
 from langchain.schema.messages import AIMessage
 from langchain.schema.messages import BaseMessage
 from langchain.schema.messages import HumanMessage
-from langchain.schema.messages import SystemMessage
 
 from onyx.agents.agent_search.dr.models import AggregatedDRContext
 from onyx.agents.agent_search.dr.models import IterationAnswer
@@ -13,9 +12,11 @@
 from onyx.agents.agent_search.shared_graph_utils.operators import (
     dedup_inference_section_list,
 )
+from onyx.configs.constants import MessageType
 from onyx.context.search.models import InferenceSection
 from onyx.context.search.models import SavedSearchDoc
 from onyx.context.search.utils import chunks_or_sections_to_search_docs
+from onyx.llm.models import PreviousMessage
 from onyx.tools.tool_implementations.web_search.web_search_tool import (
     WebSearchTool,
 )
@@ -242,38 +243,24 @@ def get_chat_history_string(chat_history: list[BaseMessage], max_messages: int)
 
 
 def get_chat_history_messages(
-    chat_history: list[BaseMessage], max_messages: int
-) -> list[SystemMessage | HumanMessage | AIMessage]:
+    chat_history: list[PreviousMessage], max_messages: int
+) -> list[HumanMessage | AIMessage]:
     """
     Get the chat history (up to max_messages) as a list of messages.
     """
-    past_messages = chat_history[-max_messages * 2 :]
-    filtered_past_messages = copy.deepcopy(past_messages)  # type: ignore
-    for past_message_number, past_message in enumerate(past_messages):
-
-        if isinstance(past_message.content, list):
-            removal_indices = []
-            for content_piece_number, content_piece in enumerate(past_message.content):
-                if (
-                    isinstance(content_piece, dict)
-                    and content_piece.get("type") != "text"
-                ):
-                    removal_indices.append(content_piece_number)
-
-            # Only rebuild the content list if there are items to remove
-            if removal_indices:
-                filtered_past_messages[past_message_number].content = [
-                    content_piece
-                    for content_piece_number, content_piece in enumerate(
-                        past_message.content
-                    )
-                    if content_piece_number not in removal_indices
-                ]
-
+    past_raw_messages = chat_history[-max_messages * 2 :]
+    filtered_past_raw_messages = []
+    for past_raw_message_number, past_raw_message in enumerate(past_raw_messages):
+        if past_raw_message.message_type == MessageType.USER:
+            filtered_past_raw_messages.append(
+                HumanMessage(content=past_raw_message.message)
+            )
         else:
-            continue
+            filtered_past_raw_messages.append(
+                AIMessage(content=past_raw_message.message)
+            )
 
-    return filtered_past_messages  # type: ignore
+    return filtered_past_raw_messages  # type: ignore
 
 
 def get_prompt_question(

backend/onyx/prompts/dr_prompts.py

@@ -1857,7 +1857,7 @@
 
 TOOL_CHOICE_WRAPPER_PROMPT = PromptTemplate(
     f"""
-Here are the tools/sub-agents and tool calls that were determined to be needed next to answer the user's question:
+Here are the tools/sub-agent calls that were determined to be needed next to answer the user's question:
 
 #TOOL CALLS
 {SEPARATOR_LINE}