Update app.py

app.py

@@ -7,7 +7,7 @@ from langchain_core.tools import tool
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langgraph.graph import StateGraph, END
 from typing import TypedDict, Annotated, Sequence
-from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_core.messages import AIMessage, HumanMessage
 import operator
 import networkx as nx
 import matplotlib.pyplot as plt
@@ -38,12 +38,11 @@ def search(query: str):
 
 tools = [search, multiply]
 tool_map = {tool.name: tool for tool in tools}
-
 model_with_tools = model.bind_tools(tools)
 
 # Define Agent State class
 class AgentState(TypedDict):
-    messages: Annotated[Sequence[BaseMessage], operator.add]
+    messages: Annotated[Sequence[HumanMessage], operator.add]
 
 # Define workflow nodes
 def invoke_model(state):
@@ -70,8 +69,9 @@ def invoke_tool(state):
         if response == "No":
             raise ValueError(f"Execution of '{selected_tool}' was canceled.")
 
+    # Invoke tool and return response as AIMessage
     response = tool_map[selected_tool].invoke(json.loads(tool_details.get("function").get("arguments")))
-    return {"messages": [BaseMessage(content=str(response))]}
+    return {"messages": [AIMessage(content=str(response))]}
 
 def router(state):
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
@@ -129,11 +129,14 @@ if st.button("Run Workflow"):
             for s in compiled_app.stream({"messages": [HumanMessage(content=prompt)]}):
                 intermediate_outputs.append(s)
 
-            # Extract and display the final response safely
+            # Extract and display the final response
             final_output = intermediate_outputs[-1]
-            messages = final_output.get('messages', [{}])[-1]
-            response_content = getattr(messages, 'content', 'No content generated.')
-            st.write("Response:", response_content)
+            messages = final_output.get('messages', [])
+            if messages:
+                response_content = messages[-1].content
+                st.write("Response:", response_content)
+            else:
+                st.write("Response: No content generated.")
         except Exception as e:
             st.error(f"Error: {e}")