Update app.py

app.py

@@ -12,7 +12,7 @@ from langgraph.prebuilt import ToolNode
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, START, END
 from langgraph.graph.message import add_messages
-from langchain.schema import HumanMessage, SystemMessage
+from langchain.schema import HumanMessage, SystemMessage, AIMessage
 # Create a ToolNode that knows about your web_search function
 import json
 from state import AgentState
@@ -61,11 +61,13 @@ def plan_node(state: AgentState) -> AgentState:
     llm_response = llm([system_msg, human_msg])
     llm_out = llm_response.content.strip()
 
-    # 4) Try to parse as a Python dict
+    # 4) Always append the LLM output as an AIMessage
+    ai_msg = AIMessage(content=llm_out)
+    new_msgs = prior_msgs.copy() + [ai_msg]
     try:
         parsed = eval(llm_out, {}, {})
         if isinstance(parsed, dict):
-            partial: AgentState = {"messages": prior_msgs.copy()}
+            partial: AgentState = {"messages": new_msgs}
             allowed = {
                 "web_search_query",
                 "ocr_path",
@@ -82,7 +84,7 @@ def plan_node(state: AgentState) -> AgentState:
 
     # 5) Fallback
     return {
-        "messages": prior_msgs.copy(),
+        "messages": new_msgs,
         "final_answer": "Sorry, I could not parse your intent."
     }
 
@@ -112,7 +114,7 @@ def finalize_node(state: AgentState) -> AgentState:
     if state.get("final_answer") is not None:
         return {"final_answer": state["final_answer"]}
 
-    # 3) Otherwise, append our “please give final answer” SystemMessage
+    # 3) Otherwise, append our "please give final answer" SystemMessage
     messages_for_llm.append(
         SystemMessage(content="Please provide the final answer now.")
     )