LangGraph: Graph-Based Agent Orchestration

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI

# 1. Define state — what data flows through the graph
class AgentState(TypedDict):
    messages: list          # Conversation history
    task: str               # The user's goal
    research_results: list  # Accumulated research
    final_answer: str       # The finished output

# 2. Define nodes — functions that transform state
def research_node(state: AgentState) -> AgentState:
    """Search for information needed to complete the task."""
    results = search_web(state["task"])
    return {**state, "research_results": state["research_results"] + [results]}

def write_answer_node(state: AgentState) -> AgentState:
    """Generate the final answer from research results."""
    llm = ChatOpenAI(model="gpt-4o")
    answer = llm.invoke(f"Task: {state['task']}\nResearch: {state['research_results']}")
    return {**state, "final_answer": answer.content}

def should_continue(state: AgentState) -> str:
    """Conditional edge — decide what to do next."""
    if len(state["research_results"]) < 2:
        return "research"  # Need more research
    return "write"         # Have enough, write answer

# 3. Build the graph
workflow = StateGraph(AgentState)
workflow.add_node("research", research_node)
workflow.add_node("write", write_answer_node)

workflow.set_entry_point("research")
workflow.add_conditional_edges(
    "research",
    should_continue,
    {"research": "research", "write": "write"}  # Maps return value → next node
)
workflow.add_edge("write", END)

# 4. Compile
graph = workflow.compile()

# 5. Run
result = graph.invoke({
    "messages": [],
    "task": "Explain the current state of quantum computing",
    "research_results": [],
    "final_answer": ""
})

from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain_community.tools import TavilySearchResults
from langchain_experimental.tools import PythonREPLTool

llm = ChatOpenAI(model="gpt-4o", temperature=0)
tools = [
    TavilySearchResults(max_results=3),
    PythonREPLTool()
]

# Creates a full ReAct agent graph internally
agent = create_react_agent(llm, tools)

# Run with streaming to see each step
config = {"recursion_limit": 25}  # Max iterations

for event in agent.stream(
    {"messages": [("human", "What is 15% of the current Bitcoin price in USD?")]},
    config=config,
    stream_mode="values"  # Stream complete state at each step
):
    last_message = event["messages"][-1]
    if hasattr(last_message, 'content') and last_message.content:
        print(f"{last_message.type}: {last_message.content[:200]}")

from typing import Annotated
from langgraph.graph.message import add_messages

class MessagesState(TypedDict):
    # add_messages: new messages are APPENDED, not replaced
    # Without this annotation, each update would overwrite the list
    messages: Annotated[list, add_messages]

# Custom state with mixed field types
class ResearchState(TypedDict):
    messages: Annotated[list, add_messages]
    sources: Annotated[list, lambda x, y: list(set(x + y))]  # Deduplicated union
    current_step: str          # Overwritten each time (no annotation)
    iteration_count: int       # Overwritten each time

from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.postgres import PostgresSaver  # For production

# Development: in-memory checkpointer
memory = MemorySaver()
agent = create_react_agent(llm, tools, checkpointer=memory)

# The thread_id identifies a specific conversation
config = {"configurable": {"thread_id": "user_alice_session_1"}}

# Turn 1
agent.invoke(
    {"messages": [("human", "My account number is ACC-123456")]},
    config
)

# Turn 2 — agent remembers the account number
result = agent.invoke(
    {"messages": [("human", "What's the balance on that account?")]},
    config  # Same thread_id = same conversation context
)

# Production: PostgreSQL checkpointer
import psycopg
conn_string = os.environ["DATABASE_URL"]
with psycopg.connect(conn_string) as conn:
    checkpointer = PostgresSaver(conn)
    checkpointer.setup()  # Creates tables if needed

agent = create_react_agent(llm, tools, checkpointer=checkpointer)

from langchain_core.messages import SystemMessage, HumanMessage, AIMessage

class ResearchAgentGraph:
    def __init__(self, llm, search_tool):
        self.llm = llm
        self.search = search_tool
        self.graph = self._build_graph()
    
    def _plan_node(self, state):
        """Generate a research plan."""
        planner = self.llm.bind(
            system="You are a research planner. Create a step-by-step research plan."
        )
        response = planner.invoke(state["messages"])
        return {"messages": state["messages"] + [response]}
    
    def _research_node(self, state):
        """Execute one research step."""
        last_message = state["messages"][-1]
        results = self.search.invoke(last_message.content)
        return {
            "messages": state["messages"] + [AIMessage(content=str(results))],
            "iteration": state.get("iteration", 0) + 1
        }
    
    def _synthesize_node(self, state):
        """Combine research into final answer."""
        synthesizer = self.llm
        response = synthesizer.invoke(state["messages"])
        return {"messages": state["messages"] + [response], "complete": True}
    
    def _route(self, state) -> str:
        if state.get("complete"):
            return END
        if state.get("iteration", 0) >= 3:
            return "synthesize"  # Force synthesis after 3 iterations
        return "research"
    
    def _build_graph(self):
        graph = StateGraph(dict)  # Using dict for flexibility
        graph.add_node("plan", self._plan_node)
        graph.add_node("research", self._research_node)
        graph.add_node("synthesize", self._synthesize_node)
        
        graph.set_entry_point("plan")
        graph.add_edge("plan", "research")
        graph.add_conditional_edges("research", self._route)
        graph.add_edge("synthesize", END)
        
        return graph.compile(checkpointer=MemorySaver())
    
    def run(self, task: str, thread_id: str) -> str:
        config = {"configurable": {"thread_id": thread_id}}
        result = self.graph.invoke(
            {"messages": [HumanMessage(content=task)], "iteration": 0},
            config
        )
        return result["messages"][-1].content

# Print ASCII visualization
graph.get_graph().print_ascii()

# Get Mermaid diagram for documentation
print(graph.get_graph().draw_mermaid())