Update app.py

app.py

@@ -45,9 +45,7 @@ def search_tavily(state):
     """ Retrieve docs from web search """
 
     # Search
-    tavily_search = TavilySearchResults(max_results=2)
-        
-    search_docs = tavily_search.invoke(state['question'])
+    tavily_search = TavilySearchResults(query=query, max_results=2)
     
     # Format
     formatted_search_docs = "\n\n---\n\n".join(
@@ -65,7 +63,7 @@ def search_wikipedia(state):
     """ Retrieve docs from wikipedia """
 
     # Search
-    search_docs = WikipediaLoader(query=state['question'], 
+    search_docs = WikipediaLoader(query=query, 
                                   load_max_docs=2).load()
     
     # Format
@@ -85,7 +83,7 @@ def router(state: MessagesState):
     return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
 
-builder = StateGraph(State)
+builder = StateGraph(MessagesState)
 builder.add_node("router", router)
 builder.add_node("tools", ToolNode([search_tavily, search_wikipedia]))
 builder.add_edge(START, "router")