langgraph_agent.py

# langgraph_agent.py
import os
from langgraph.graph import StateGraph, END, START, MessagesState
from langgraph.prebuilt import ToolNode
from langchain.agents import Tool
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.prebuilt import tools_condition

from tools.langgraph_tools import wiki_search, web_search

from dotenv import load_dotenv
load_dotenv()

api_key = os.getenv("OPENAI_API_KEY")

# Define some tools
def echo_tool(input: str) -> str:
    return f"Echo: {input}"

tools = [
    wiki_search,
    # web_search,
]

# LangChain-compatible LLM
llm = ChatOpenAI(api_key=api_key, model="gpt-4o")
llm_with_tools = llm.bind_tools(tools)

# Tool executor node
tool_node = ToolNode(tools)

# Node
def assistant(state: MessagesState):
    """Assistant node"""
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

def get_agent():
    graph = StateGraph(MessagesState)
    graph.add_node("assistant", assistant)
    graph.add_node("tools", tool_node)
    graph.add_edge(START, "assistant")
    graph.add_edge("assistant", "tools")
    graph.add_conditional_edges(
        "assistant",
        tools_condition,
    )
    # graph.add_edge("tools", "assistant")
    graph.add_edge("assistant", END)
    return graph.compile()

# test
if __name__ == "__main__":
    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
    # Build the graph
    agent = get_agent()
    # Run the graph
    messages = [HumanMessage(content=question)]
    messages = agent.invoke({"messages": messages}, {"recursion_limit": 5})
    for m in messages["messages"]:
        m.pretty_print()