Update app.py

app.py

@@ -5,12 +5,12 @@ import inspect
 import pandas as pd
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langgraph.graph import StateGraph, MessagesState, START
+from langgraph.prebuilt import ToolNode
+from langgraph.prebuilt import tools_condition
+from langchain_core.tools import tool
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_community.document_loaders import WikipediaLoader
 from langchain_community.tools import TavilySearchResults
-import operator
-from typing import Annotated
-from typing_extensions import TypedDict
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -33,20 +33,13 @@ class BasicAgent:
         print("BasicAgent initialized.")
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        initial_state = {
-        "question": "What is the capital of France?",
-        "context": []
-        }
-        final_state = self.graph.invoke(initial_state)
-        answer = final_state["answer"]
+        messages = [SystemMessage(content=SYSTEM_MESSAGE)] + [HumanMessage(content=f"Answer the question: {question}")]
+        messages = self.graph.invoke({"messages": messages})
+        answer = messages['messages'][-1].content
         print(f"Agent returning fixed answer: {answer}")
         return answer
 
-class State(TypedDict):
-    question: str
-    answer: str
-    context: Annotated[list, operator.add]
-    
+@tool    
 def search_tavily(state):
     
     """ Retrieve docs from web search """
@@ -66,6 +59,7 @@ def search_tavily(state):
 
     return {"context": [formatted_search_docs]}
 
+@tool    
 def search_wikipedia(state):
     
     """ Retrieve docs from wikipedia """
@@ -84,34 +78,19 @@ def search_wikipedia(state):
 
     return {"context": [formatted_search_docs]}
 
-def generate_answer(state):
-
-    """Node to give answer to the question"""
-
-    context = state["context"]
-    question = state["question"]
-
-    additional_context_template = """Here are some contexts about the question you can use if you find it helpful: {context}"""
-    additional_context = additional_context_template.format(context=context)
-    final_instruction = SYSTEM_MESSAGE + additional_context
+llm_with_tools = llm.bind_tools([search_tavily, search_wikipedia])
 
-    #answer
-    answer = llm.invoke([SystemMessage(content=final_instruction)] + [HumanMessage(content=f"Answer the question: {question}")])
-
-    # Append it to state
-    return {"answer": answer}
+def router(state: MessagesState):
+    """Router of the graph"""
+    return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
 
 builder = StateGraph(State)
-
-builder.add_node("search_tavily",search_tavily)
-builder.add_node("search_wikipedia", search_wikipedia)
-builder.add_node("generate_answer", generate_answer)
-
-builder.add_edge(START, "search_wikipedia")
-builder.add_edge(START, "search_tavily")
-builder.add_edge("search_wikipedia", "generate_answer")
-builder.add_edge("search_tavily", "generate_answer")
+builder.add_node("router", router)
+builder.add_node("tools", ToolNode([search_tavily, search_wikipedia]))
+builder.add_edge(START, "router")
+builder.add_conditional_edges("router", tools_condition)
+builder.add_edge("tools", "router")
 graph = builder.compile()