From Solo Agent to Team Player: Architecting Multi-Agent Systems with LangGraph

from typing import TypedDict, List, Optional
from langchain_core.messages import BaseMessage
 
class AgentState(TypedDict):
    """
    Represents the state of our multi-agent system.
    Shared context accessible and modifiable by all agents.
    """
    query: str                       # The initial user query/topic
    research_results: List[str]      # Accumulated research findings
    draft: Optional[str]             # Current draft of the summary
    review_comments: List[str]       # Feedback from the reviewer
    iterations: int                  # Number of review-refine cycles
    max_iterations: int              # Maximum allowed review-refine cycles
    next_action: Optional[str]       # Supervisor's decision on the next agent to invoke
    messages: List[BaseMessage]      # For conversational history if needed (e.g., LangGraph's default)

import os
from openai import OpenAI
from typing import Callable, Literal
from langchain_core.messages import HumanMessage, SystemMessage, FunctionMessage
from langgraph.graph import StateGraph, END
 
# --- LLM Client Setup ---
# Use raw OpenAI client for direct API interaction, avoiding LangChain's LLM wrappers.
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
 
def call_llm(prompt: str, model: str = "gpt-4o", temperature: float = 0.7) -> str:
    """Helper to make direct LLM calls."""
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=temperature,
    )
    return response.choices[0].message.content
 
# --- Agent Definitions ---
 
class Agents:
    """Contains our specialized agent functions."""
 
    def __init__(self, llm_model: str = "gpt-4o"):
        self.llm_model = llm_model
 
    def supervisor_agent(self, state: AgentState) -> AgentState:
        """
        The orchestrator. Decides which worker agent to invoke next or if the task is complete.
        Returns a state dict with 'next_action'.
        """
        current_draft_status = "no draft yet"
        if state.get("draft"):
            current_draft_status = "draft exists"
 
        review_status = "no review comments"
        if state.get("review_comments"):
            review_status = f"{len(state['review_comments'])} review comments exist"
 
        supervisor_prompt = f"""
        You are the Supervisor. Your goal is to manage a research and content generation team.
        The current query is: "{state['query']}"
 
        Current Status:
        - Research Results: {len(state['research_results'])} pieces of information gathered.
        - Draft Status: {current_draft_status}
        - Review Status: {review_status}
        - Iterations: {state['iterations']}/{state['max_iterations']}
 
        Based on the current state, decide the next logical step.
        Choose one of the following actions: 'research', 'draft', 'review', 'revise', 'finish'.
 
        Rules:
        1. If no research has been done, or existing research is insufficient, choose 'research'.
        2. If research is sufficient and no draft exists, choose 'draft'.
        3. If a draft exists and has not been reviewed (or needs further review), choose 'review'.
        4. If review comments exist and iterations < max_iterations, choose 'revise'.
        5. If a draft has been reviewed, all comments addressed (or max_iterations reached), choose 'finish'.
 
        Your output MUST be a single word, one of: research, draft, review, revise, finish.
        """
        print(f"--- Supervisor Thinking ---")
        action = call_llm(supervisor_prompt, model=self.llm_model, temperature=0.1).strip().lower()
        print(f"Supervisor decided: {action}")
        return {"next_action": action}
 
    def researcher_agent(self, state: AgentState) -> AgentState:
        """
        Gathers information based on the query.
        (For this example, we'll simulate web search with an LLM call.)
        """
        print(f"--- Researcher Working ---")
        research_prompt = f"""
        You are a highly skilled researcher. Find 3-5 key facts or pieces of information
        about the following topic: "{state['query']}".
        Focus on technical aspects and provide concise bullet points.
        If previous research results exist, try to expand on them or find complementary information.
        Previous results: {state['research_results']}
        """
        new_research = call_llm(research_prompt, model=self.llm_model).strip()
        
        updated_results = state.get("research_results", []) + [new_research]
        print(f"Researcher found: {new_research[:100]}...")
        return {"research_results": updated_results}
 
    def writer_agent(self, state: AgentState) -> AgentState:
        """
        Drafts the technical summary based on research results.
        """
        print(f"--- Writer Working ---")
        research_summary = "\n".join(state["research_results"])
        writer_prompt = f"""
        You are a technical writer. Draft a concise, informative summary (approx. 200-300 words)
        on the topic: "{state['query']}".
        Use the following research results:
        ---
        {research_summary}
        ---
        Focus on clarity, accuracy, and technical detail appropriate for a developer audience.
        """
        draft = call_llm(writer_prompt, model=self.llm_model).strip()
        print(f"Writer drafted: {draft[:100]}...")
        return {"draft": draft}
 
    def reviewer_agent(self, state: AgentState) -> AgentState:
        """
        Reviews the draft and provides constructive criticism.
        """
        print(f"--- Reviewer Working ---")
        draft = state["draft"]
        reviewer_prompt = f"""
        You are a critical technical reviewer. Evaluate the following draft for accuracy,
        clarity, completeness, and adherence to the topic "{state['query']}".
        Provide constructive feedback, suggesting specific improvements. If the draft is perfect,
        state "No specific changes needed."
        Draft:
        ---
        {draft}
        ---
        """
        comments = call_llm(reviewer_prompt, model=self.llm_model).strip()
        
        updated_comments = state.get("review_comments", []) + [comments]
        print(f"Reviewer commented: {comments[:100]}...")
        return {"review_comments": updated_comments}
 
    def reviser_agent(self, state: AgentState) -> AgentState:
        """
        Revises the draft based on reviewer comments.
        Increments the iteration count.
        """
        print(f"--- Reviser Working ---")
        original_draft = state["draft"]
        review_feedback = "\n".join(state["review_comments"])
        reviser_prompt = f"""
        You are a meticulous reviser. Improve the following draft based on the provided
        review feedback. Ensure all comments are addressed while maintaining the
        original intent and technical accuracy for the topic "{state['query']}".
        Original Draft:
        ---
        {original_draft}
        ---
        Review Feedback:
        ---
        {review_feedback}
        ---
        Provide the revised draft.
        """
        revised_draft = call_llm(reviser_prompt, model=self.llm_model).strip()
        
        # Clear review comments for the next cycle and increment iteration
        new_iterations = state.get("iterations", 0) + 1
        print(f"Reviser revised. Iteration: {new_iterations}")
        return {"draft": revised_draft, "review_comments": [], "iterations": new_iterations}
 

# Initialize our agents
team_agents = Agents()
 
# Build the graph
workflow = StateGraph(AgentState)
 
# Add nodes for each agent function
workflow.add_node("supervisor", team_agents.supervisor_agent)
workflow.add_node("researcher", team_agents.researcher_agent)
workflow.add_node("writer", team_agents.writer_agent)
workflow.add_node("reviewer", team_agents.reviewer_agent)
workflow.add_node("reviser", team_agents.reviser_agent)
 
# Set the entry point - always start with the supervisor
workflow.set_entry_point("supervisor")
 
# Define conditional edges from the supervisor
workflow.add_conditional_edges(
    "supervisor",
    lambda state: state["next_action"], # The supervisor's decision dictates the next node
    {
        "research": "researcher",
        "draft": "writer",
        "review": "reviewer",
        "revise": "reviser",
        "finish": END # If supervisor says 'finish', the graph ends
    }
)
 
# Define regular edges for worker agents to return to the supervisor
workflow.add_edge("researcher", "supervisor")
workflow.add_edge("writer", "supervisor")
workflow.add_edge("reviewer", "supervisor")
workflow.add_edge("reviser", "supervisor")
 
# Compile the graph
app = workflow.compile()
 
# Optional: Visualize the graph (requires pygraphviz or pydot)
# from IPython.display import Image, display
# display(Image(app.get_graph().draw_png()))

# --- Example Execution ---
initial_state: AgentState = {
    "query": "Explain transformer models in AI.",
    "research_results": [],
    "draft": None,
    "review_comments": [],
    "iterations": 0,
    "max_iterations": 2, # Allow up to 2 review/revise cycles
    "next_action": None,
    "messages": [] # LangGraph internally uses this for conversational state sometimes
}
 
print(f"\n--- Starting Research for: {initial_state['query']} ---")
final_state = {}
for s in app.stream(initial_state):
    print(f"Current State: {list(s.keys())[0]} - {s}")
    final_state.update(s)
    print("---")
 
print("\n--- Final Output ---")
print(f"Topic: {final_state['supervisor']['query']}")
print(f"Final Draft:\n{final_state['supervisor']['draft']}")
print(f"Total Iterations: {final_state['supervisor']['iterations']}")