diff --git a/main.py b/main.py index 03ee0ef..770317e 100644 --- a/main.py +++ b/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.