update main.py

main.py

@@ -1,7 +1,7 @@
 """
 Main entry point for the agent with memory and human‑in‑the‑loop confirmation.
 
-The agent is built on top of LangGraph's `create_agent` API.  It uses a
+The agent is built on top of LangChain's `create_agent` API.  It uses a
 `MemorySaver` checkpoint to keep conversation history across calls, and it
 is configured with `interrupt_before=["tools"]` so that the agent pauses just
 before invoking any tool.  The pause allows us to ask the user for explicit
@@ -30,8 +30,8 @@ from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage, SystemMessage
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.types import Command
-from langgraph.graph import StateGraph
-from langgraph.graph.message import add_messages
+from langchain.agents import create_agent
+from langchain.tools import tool
 from rich.console import Console
 
 # ---------------------------------------------------------------------------
@@ -49,6 +49,7 @@ llm = ChatOpenAI(
 # ---------------------------------------------------------------------------
 # Simple tool – in a real scenario replace with an actual API call.
 # ---------------------------------------------------------------------------
+@tool
 async def get_price(query: Dict[str, Any]) -> str:
     """Pretend to fetch a price for a product.
 
@@ -68,49 +69,23 @@ async def get_price(query: Dict[str, Any]) -> str:
     return f"The price of {product} is 42.00 {currency}."
 
 # ---------------------------------------------------------------------------
-# Agent definition
+# Agent definition using create_agent (LangChain)
 # ---------------------------------------------------------------------------
 memory = MemorySaver()
 
-agent = StateGraph(
-    state_schema=dict(messages=list, next=tuple)
+agent = create_agent(
+    llm=llm,
+    tools=[get_price],
+    system_prompt="You are a helpful assistant that can query prices.",
+    checkpointer=memory,
+    interrupt_before=["tools"],  # pause before any tool call
 )
 
-# Node that simply forwards the messages to the LLM.
-async def llm_node(state: Dict[str, Any]) -> Tuple[Dict[str, Any], str]:
-    # The LLM expects a list of messages; we pass the current history.
-    response = await llm.ainvoke(state["messages"])
-    return {"messages": state["messages"] + [response]}, "next"
-
-# Node that handles tool calls – for this example we only have get_price.
-async def tool_node(state: Dict[str, Any]) -> Tuple[Dict[str, Any], str]:
-    # The last message should contain a tool call.
-    last_msg = state["messages"][-1]
-    if not hasattr(last_msg, "tool_calls") or not last_msg.tool_calls:
-        return state, "next"
-
-    tool_call = last_msg.tool_calls[0]
-    name = tool_call.name
-    args = json.loads(tool_call.args)
-    if name == "get_price":
-        result = await get_price(args)
-        # Append the tool output as a new message.
-        state["messages"].append(
-            HumanMessage(content=f"Tool {name} returned: {result}")
-        )
-    return state, "next"
-
-agent.add_node("llm", llm_node)
-agent.add_node("tool", tool_node)
-agent.set_entry_point("llm")
-agent.add_edge("llm", "tool")
-agent.add_edge("tool", "llm")
-
-# Compile the graph with a MemorySaver checkpoint.
+# Compile the agent into a graph with checkpointing.
 graph = agent.compile(checkpointer=memory, interrupt_before=["tools"])
 
 # ---------------------------------------------------------------------------
-# Helper to run the agent with human‑in‑the‑loop confirmation.
+# Helper to run the agent and pause before each tool call.
 # ---------------------------------------------------------------------------
 async def ask_and_run(user_input: Dict[str, Any], config: Dict[str, Any]):
     """Run the agent and pause before each tool call.
@@ -157,10 +132,11 @@ if __name__ == "__main__":
     console.print("[bold green]Welcome to the agent demo![/]")
     console.print("Type 'exit' to quit.")
 
+    import asyncio
+
     # Example 1 – simple chat (no tool call).
     console.print("\n[underline]Example 1: Simple question[/]")
     user_msg = {"messages": [HumanMessage(content="What is the capital of France?")]} 
-    import asyncio
     asyncio.run(ask_and_run(user_msg, config))
 
     # Example 2 – tool call with confirmation.