Update app.py

app.py

@@ -17,7 +17,6 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.tools.retriever import create_retriever_tool
 from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
 from langgraph.graph import StateGraph, START, END
-from langgraph.graph.state import MessagesState
 from langgraph.prebuilt import ToolNode, tools_condition
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import CompiledStateGraph
@@ -53,7 +52,6 @@ def space_events_agent():
             category = (event.get('categories', [{}])[0]).get('title', 'Unknown Category')
             geometries = event.get('geometries', [])
             coordinates = geometries[0].get('coordinates', 'Unknown Location') if geometries else 'Unknown Location'
-            # Include additional detailed formatting for each event
             event_details.append(
                 f"• **Event:** {event.get('title', 'N/A')}\n"
                 f"  **Type:** {category}\n"
@@ -238,12 +236,12 @@ sys_msg = SystemMessage(content='''This system is designed to provide real-time
 When responding to user queries, the assistant should provide comprehensive, detailed, and accurate explanations. Visual responses should include image URLs that the interface can use to display pictures. All responses should strive for clarity and depth in explanation.
 ''')
 
-# Define the assistant node function
-def assistant(state: MessagesState) -> MessagesState:
+# Define the assistant node function; we use a plain dictionary as the state type
+def assistant(state: dict) -> dict:
     return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"][-10:])]}
 
-# Build the LangGraph state graph
-builder: StateGraph = StateGraph(MessagesState)
+# Build the LangGraph state graph using dict as the state type
+builder: StateGraph = StateGraph(dict)
 builder.add_node("assistant", assistant)
 builder.add_node("tools", ToolNode(tools))
 builder.add_edge(START, "assistant")
@@ -282,7 +280,7 @@ if user_input := st.chat_input("Ask about space news, ISS location, astronomy im
         # Check for image URL markers in a case-insensitive way
         image_url = None
         if "image url:" in assistant_response.lower():
-            # Use regex to extract the first URL after "image url:"
+            # Use regex to extract the first URL after "Image URL:"
             match = re.search(r"[Ii]mage [Uu][Rr][Ll]:\s*(\S+)", assistant_response)
             if match:
                 image_url = match.group(1).strip()