include system message

Author: evan-danswer

Diff for: backend/onyx/chat/prompt_builder/answer_prompt_builder.py

@@ -155,6 +155,19 @@ def get_user_message_content(self) -> str:
         query, _ = message_to_prompt_and_imgs(self.user_message_and_token_cnt[0])
         return query
 
+    def get_message_history(self) -> list[PreviousMessage]:
+        """
+        Get the message history as a list of PreviousMessage objects.
+        """
+        ret = []
+        if self.system_message_and_token_cnt:
+            tmp = PreviousMessage.from_langchain_msg(*self.system_message_and_token_cnt)
+            ret.append(tmp)
+        for i, msg in enumerate(self.message_history):
+            tmp = PreviousMessage.from_langchain_msg(msg, self.history_token_cnts[i])
+            ret.append(tmp)
+        return ret
+
     def build(self) -> list[BaseMessage]:
         if not self.user_message_and_token_cnt:
             raise ValueError("User message must be set before building prompt")

Diff for: backend/onyx/chat/tool_handling/tool_response_handler.py

@@ -163,8 +163,10 @@ def get_tool_call_for_non_tool_calling_llm_impl(
     llm: LLM,
 ) -> tuple[Tool, dict] | None:
     user_query = prompt_builder.raw_user_query
+    history = prompt_builder.raw_message_history
     if isinstance(prompt_builder, AnswerPromptBuilder):
         user_query = prompt_builder.get_user_message_content()
+        history = prompt_builder.get_message_history()
 
     if force_use_tool.force_use:
         # if we are forcing a tool, we don't need to check which tools to run
@@ -175,7 +177,7 @@ def get_tool_call_for_non_tool_calling_llm_impl(
             if force_use_tool.args is not None
             else tool.get_args_for_non_tool_calling_llm(
                 query=user_query,
-                history=prompt_builder.raw_message_history,
+                history=history,
                 llm=llm,
                 force_run=True,
             )
@@ -193,7 +195,7 @@ def get_tool_call_for_non_tool_calling_llm_impl(
         tool_options = check_which_tools_should_run_for_non_tool_calling_llm(
             tools=tools,
             query=user_query,
-            history=prompt_builder.raw_message_history,
+            history=history,
             llm=llm,
         )
 
@@ -210,7 +212,7 @@ def get_tool_call_for_non_tool_calling_llm_impl(
         chosen_tool_and_args = (
             select_single_tool_for_non_tool_calling_llm(
                 tools_and_args=available_tools_and_args,
-                history=prompt_builder.raw_message_history,
+                history=history,
                 query=user_query,
                 llm=llm,
             )

Diff for: backend/onyx/llm/models.py

@@ -9,6 +9,7 @@
 from onyx.configs.constants import MessageType
 from onyx.file_store.models import InMemoryChatFile
 from onyx.llm.utils import build_content_with_imgs
+from onyx.llm.utils import message_to_string
 from onyx.tools.models import ToolCallFinalResult
 
 if TYPE_CHECKING:
@@ -59,3 +60,22 @@ def to_langchain_msg(self) -> BaseMessage:
             return AIMessage(content=content)
         else:
             return SystemMessage(content=content)
+
+    @classmethod
+    def from_langchain_msg(
+        cls, msg: BaseMessage, token_count: int
+    ) -> "PreviousMessage":
+        message_type = MessageType.SYSTEM
+        if isinstance(msg, HumanMessage):
+            message_type = MessageType.USER
+        elif isinstance(msg, AIMessage):
+            message_type = MessageType.ASSISTANT
+        message = message_to_string(msg)
+        return cls(
+            message=message,
+            token_count=token_count,
+            message_type=message_type,
+            files=[],
+            tool_call=None,
+            refined_answer_improvement=None,
+        )