Convert consecutive ToolMessages to AIMessage

app.py

@@ -26,7 +26,9 @@ if torch.cuda.is_available():
     ckpt_dir = snapshot_download(model_id, local_dir_use_symlinks=False)
     print(f"Using checkpoints from {ckpt_dir}")
     print(f"Downloading checkpoints for {embedding_model_id}...")
-    embedding_ckpt_dir = snapshot_download(embedding_model_id, local_dir_use_symlinks=False)
+    embedding_ckpt_dir = snapshot_download(
+        embedding_model_id, local_dir_use_symlinks=False
+    )
     print(f"Using embedding checkpoints from {embedding_ckpt_dir}")
 else:
     ckpt_dir = None
@@ -173,7 +175,7 @@ def run_workflow(input, history, compute_mode, thread_id, session_hash):
                 count += 1
                 # Get the retrieved emails as a list
                 email_list = message.content.replace(
-                    "### Retrieved Emails:\n\n\n\n", ""
+                    "### Retrieved Emails:\n\n", ""
                 ).split("--- --- --- --- Next Email --- --- --- ---\n\n")
                 # Get the list of source files (e.g. R-help/2024-December.txt) for retrieved emails
                 month_list = [text.splitlines()[0] for text in email_list]
@@ -196,7 +198,7 @@ def run_workflow(input, history, compute_mode, thread_id, session_hash):
                 # Format the retrieved emails with Tool Call heading
                 retrieved_emails.append(
                     message.content.replace(
-                        "### Retrieved Emails:\n\n\n\n",
+                        "### Retrieved Emails:\n\n",
                         f"### ### ### ### Tool Call {count} ### ### ### ###\n\n",
                     )
                 )

graph.py

@@ -44,24 +44,45 @@ def print_message_summaries(messages, header):
 
 
 def normalize_messages(messages):
-    """Normalize messages to sequence of types expected by chat templates"""
+    """Normalize messages to sequence of types expected by chat models"""
     # Copy the most recent HumanMessage to the end
-    # (avoids SmolLM and Qwen ValueError: Last message must be a HumanMessage!)
+    # - Avoids SmolLM and Qwen ValueError: Last message must be a HumanMessage!
     if not type(messages[-1]) is HumanMessage:
         for msg in reversed(messages):
             if type(msg) is HumanMessage:
                 messages.append(msg)
                 break
-    # Convert tool output (ToolMessage) to AIMessage
-    # (avoids SmolLM and Qwen ValueError: Unknown message type: <class 'langchain_core.messages.tool.ToolMessage'>)
-    messages = [
-        AIMessage(msg.content) if type(msg) is ToolMessage else msg for msg in messages
-    ]
+
+    # Convert tool output (one or more consecutive ToolMessages) to AIMessage
+    # - Avoids SmolLM and Qwen ValueError: Unknown message type: <class 'langchain_core.messages.tool.ToolMessage'>
+    processed_messages = []
+    i = 0
+    while i < len(messages):
+        msg = messages[i]
+
+        if type(msg) is ToolMessage:
+            # Collect consecutive ToolMessages
+            tool_messages = []
+            count = 1
+            while i < len(messages) and type(messages[i]) is ToolMessage:
+                tool_msg = messages[i]
+                formatted_content = f"## Tool Call {count}\n\n{tool_msg.content}"
+                tool_messages.append(formatted_content)
+                count += 1
+                i += 1
+
+            # Combine all tool message contents into a single AIMessage
+            combined_content = "\n\n".join(tool_messages)
+            processed_messages.append(AIMessage(combined_content))
+        else:
+            processed_messages.append(msg)
+            i += 1
+
     # Delete tool call (AIMessage)
-    # (avoids Gemma TemplateError: Conversation roles must alternate user/assistant/user/assistant/...)
+    # - Avoids Gemma TemplateError: Conversation roles must alternate user/assistant/user/assistant/...
     messages = [
         msg
-        for msg in messages
+        for msg in processed_messages
         if not hasattr(msg, "tool_calls")
         or (hasattr(msg, "tool_calls") and not msg.tool_calls)
     ]
@@ -168,12 +189,12 @@ def BuildGraph(
             search_query = " ".join([search_query, start_year, end_year])
         retrieved_docs = retriever.invoke(search_query)
         serialized = "\n\n--- --- --- --- Next Email --- --- --- ---".join(
-            # source key has file names (e.g. R-help/2024-December.txt), useful for retrieval and reporting
+            # Add file name (e.g. R-help/2024-December.txt) from source key
             "\n\n" + doc.metadata["source"] + doc.page_content
             for doc in retrieved_docs
         )
         retrieved_emails = (
-            "### Retrieved Emails:\n\n" + serialized
+            "### Retrieved Emails:" + serialized
             if serialized
             else "### No emails were retrieved"
         )

main.py

@@ -207,15 +207,14 @@ def RunChain(
 
     # Create a prompt template
     system_template = ChatPromptTemplate.from_messages([SystemMessage(system_prompt)])
+    # NOTE: Each new email starts with \n\n\nFrom, so we don't need newlines after Retrieved Emails:
     human_template = ChatPromptTemplate.from_template(
         """"
         ### Question:
 
         {question}
 
-        ### Retrieved Emails:
-
-        {context}
+        ### Retrieved Emails:{context}
         """
     )
     prompt_template = system_template + human_template

requirements.txt

@@ -9,6 +9,10 @@ chromadb==1.0.13
 # NOTE: Gemma 3 with transformers==4.54.0 gives:
 #   ValueError: Max cache length is not consistent across layers
 transformers==4.51.3
+# Required by langchain-huggingface
+sentence-transformers==5.0.0
+# For snapshot_download
+huggingface-hub==0.34.3
 
 # Langchain packages
 langchain==0.3.26
@@ -23,8 +27,6 @@ langgraph-sdk==0.1.72
 langgraph-prebuilt==0.5.2
 langgraph-checkpoint==2.1.0
 
-# Required by langchain-huggingface
-sentence-transformers==5.0.0
 # Required by Nomic embeddings
 einops==0.8.1