agent graph improvements

agent.py

@@ -1,27 +1,48 @@
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_core.messages import HumanMessage, SystemMessage
-from tools import *
+from langchain_core.messages import HumanMessage, SystemMessage, AnyMessage
+from typing import TypedDict, Annotated
+from langgraph.graph.message import add_messages
+from langgraph.prebuilt import ToolNode
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import tools_condition
+from tools import tools
 
 model = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0)
+model_with_tools = model.bind_tools(tools)
 
 system_template = "You are a general AI assistant. I will ask you a question. Your final answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."
-messages = [
-    SystemMessage(system_template),
-    HumanMessage("I'm making a grocery list for my mom, but she's a professor of botany and she's a real stickler when it comes to categorizing things. I need to add different foods to different categories on the grocery list, but if I make a mistake, she won't buy anything inserted in the wrong category. Here's the list I have so far:\n\nmilk, eggs, flour, whole bean coffee, Oreos, sweet potatoes, fresh basil, plums, green beans, rice, corn, bell pepper, whole allspice, acorns, broccoli, celery, zucchini, lettuce, peanuts\n\nI need to make headings for the fruits and vegetables. Could you please create a list of just the vegetables from my list? If you could do that, then I can figure out how to categorize the rest of the list into the appropriate categories. But remember that my mom is a real stickler, so make sure that no botanical fruits end up on the vegetable list, or she won't get them when she's at the store. Please alphabetize the list of vegetables, and place each item in a comma separated list."),
-]
 
-tools = [web_search]
+class AgentState(TypedDict):
+    messages: Annotated[list[AnyMessage], add_messages]
 
-model_with_tools = model.bind_tools(tools)
+def assistant(state: AgentState):
+    return {
+        "messages": [model_with_tools.invoke(state["messages"])],
+    }
+
+## The graph
+builder = StateGraph(AgentState)
+
+# Define nodes: these do the work
+builder.add_node("assistant", assistant)
+builder.add_node("tools", ToolNode(tools))
+
+# Define edges: these determine how the control flow moves
+builder.add_edge(START, "assistant")
+builder.add_conditional_edges(
+    "assistant",
+    # If the latest message requires a tool, route to tools
+    # Otherwise, provide a direct response
+    tools_condition,
+)
+builder.add_edge("tools", "assistant")
+chat = builder.compile()
+
+config = {"configurable": {"thread_id": "1"}}
 
 def invoke_agent(query: str) -> str:
     """
     Invoke the agent with a given question.
     """
-    messages = [
-        SystemMessage(system_template),
-        HumanMessage(query),
-    ]
-
-    response = model.invoke(messages)
-    return response.text()
+    response = chat.invoke({"messages": [HumanMessage(content=query)]})
+    return response['messages'][-1].content

requirements.txt

@@ -1,5 +1,6 @@
 gradio
 requests
+langgraph
 langchain
 langchain-community
 langchain-google-genai

tools.py

@@ -1,3 +1,7 @@
-from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_community.tools import DuckDuckGoSearchResults
 
-web_search = DuckDuckGoSearchRun()
+web_search = DuckDuckGoSearchResults()
+
+tools = [
+    web_search
+]