Pythonlanggraph_agent.py
Copy
1
from typing import Annotated, TypedDict
2
from langgraph.graph import StateGraph, START, END
3
from langgraph.graph.message import add_messages
4
from langgraph.prebuilt import ToolNode
5
from langchain_openai import ChatOpenAI
6
from langchain_core.tools import tool
7
 
8
llm = ChatOpenAI(model="gpt-4o", temperature=0)
9
 
10
class State(TypedDict):
11
    messages: Annotated[list, add_messages]
12
 
13
@tool
14
def search(query: str) -> str:
15
    """Search for information about a topic."""
16
    return f"Search results for '{query}': Top 3 articles found with relevant content."
17
 
18
@tool  
19
def calculate(expression: str) -> str:
20
    """Evaluate a math expression. Only numbers and +,-,*,/,(,) allowed."""
21
    try:
22
        allowed = set("0123456789+-*/(). ")
23
        if any(c not in allowed for c in expression):
24
            return f"Error: invalid characters"
25
        return str(eval(expression))
26
    except Exception as e:
27
        return f"Error: {e}"
28
 
29
tools = [search, calculate]
30
llm_with_tools = llm.bind_tools(tools)
31
tool_node = ToolNode(tools)
32
 
33
def agent_node(state: State) -> State:
34
    response = llm_with_tools.invoke(state["messages"])
35
    return {"messages": [response]}
36
 
37
def router(state: State) -> str:
38
    last = state["messages"][-1]
39
    if hasattr(last, "tool_calls") and last.tool_calls:
40
        return "tools"
41
    return END
42
 
43
graph = StateGraph(State)
44
graph.add_node("agent", agent_node)
45
graph.add_node("tools", tool_node)
46
graph.add_edge(START, "agent")
47
graph.add_conditional_edges("agent", router, {"tools": "tools", END: END})
48
graph.add_edge("tools", "agent")
49
 
50
app = graph.compile()
51
 
52
if __name__ == "__main__":
53
    result = app.invoke({"messages": [("user", "What is (sqrt(144) + 25) * 2?")]})
54
    print(result["messages"][-1].content)
55