Taming Autonomous Agents:Multi-Agent Workflows with LangGraph

from typing import TypedDict, Annotated, Optional
from langgraph.graph.message import add_messages

class ResearchState(TypedDict):
    """
    Shared state that flows through every node in the graph.
    Each field is written by a specific node and read by downstream nodes.
    """
    # Input
    topic: str
    
    # Written by: Researcher
    research_findings: Optional[str]
    sources: Optional[list[str]]
    
    # Written by: Writer
    draft: Optional[str]
    
    # Written by: Verifier
    verification_status: Optional[str]  # "approved" | "needs_revision"
    rejection_reason: Optional[str]
    
    # Loop control — prevent infinite cycles
    revision_count: int
    max_revisions: int
    
    # Final output
    final_document: Optional[str]
    error: Optional[str]

from langchain_openai import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import SystemMessage, HumanMessage
from state import ResearchState

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
search_tool = TavilySearchResults(max_results=5)

RESEARCHER_PROMPT = """You are a precise research analyst. Your job is to gather factual 
information about the given topic. Search for current, accurate data. 
Return structured findings with source URLs. Do NOT summarize or write prose — 
only organize the raw information you find. Flag anything that seems uncertain."""

def researcher_node(state: ResearchState) -> dict:
    """
    Searches for information on the topic and stores raw findings in state.
    """
    print(f"[Researcher] Researching: {state['topic']}")
    
    # Perform web searches
    search_results = search_tool.invoke(state["topic"])
    
    # Format search results for the LLM to organize
    raw_results = "\n\n".join([
        f"Source: {r['url']}\nContent: {r['content']}"
        for r in search_results
    ])
    
    # Have the LLM organize the findings (not summarize — organize)
    messages = [
        SystemMessage(content=RESEARCHER_PROMPT),
        HumanMessage(content=f"Topic: {state['topic']}\n\nSearch Results:\n{raw_results}\n\nOrganize these findings.")
    ]
    
    response = llm.invoke(messages)
    sources = [r['url'] for r in search_results]
    
    print(f"[Researcher] Found {len(sources)} sources")
    
    return {
        "research_findings": response.content,
        "sources": sources
    }

from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage
from state import ResearchState

llm = ChatOpenAI(model="gpt-4o", temperature=0.3)

WRITER_PROMPT = """You are an expert technical writer. Write a clear, accurate, 
well-structured document based ONLY on the research findings provided. 
Do not add information that is not in the research. 
Do not speculate. If the research is incomplete on a point, say so explicitly.
Write for a technical founder audience — sophisticated but time-constrained."""

def writer_node(state: ResearchState) -> dict:
    """
    Drafts a document from research findings.
    If this is a revision, incorporates the rejection feedback.
    """
    revision_count = state.get("revision_count", 0)
    rejection_reason = state.get("rejection_reason", "")
    
    print(f"[Writer] Drafting document (revision {revision_count})")
    
    # Build the prompt — include rejection feedback for revisions
    user_content = f"""Research Findings:
{state['research_findings']}

Sources Used: {', '.join(state.get('sources', []))}
"""
    
    if rejection_reason:
        user_content += f"""

IMPORTANT: This is revision #{revision_count}. 
The previous draft was rejected for this reason: {rejection_reason}
You MUST address this specific feedback in your revision."""
    
    messages = [
        SystemMessage(content=WRITER_PROMPT),
        HumanMessage(content=user_content)
    ]
    
    response = llm.invoke(messages)
    
    return {
        "draft": response.content,
        "revision_count": revision_count + 1
    }

from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage
import json
from state import ResearchState

# Use temperature=0 for verification — we want consistent, deterministic assessment
llm = ChatOpenAI(model="gpt-4o", temperature=0)

VERIFIER_PROMPT = """You are a rigorous fact-checker and technical editor.
Review the drafted document against the provided research findings.

Check for:
1. Factual accuracy — are all claims supported by the research?
2. No hallucinations — is any information NOT in the research sources?
3. Logical consistency — are there any contradictions?
4. Completeness — are key aspects of the topic missing?

Respond with a JSON object:
{"status": "approved"|"needs_revision", "reason": "...", "confidence": 0.0-1.0}

Be strict. If you have any doubt, reject with a specific reason."""

def verifier_node(state: ResearchState) -> dict:
    """Reviews draft against research findings. Returns approval or rejection."""
    messages = [
        SystemMessage(content=VERIFIER_PROMPT),
        HumanMessage(content=f"""Research:\n{state["research_findings"]}\n\nDraft:\n{state["draft"]}\n\nAssess as JSON.""")
    ]
    response = llm.invoke(messages)

    try:
        # Strip markdown fences if the LLM wraps the JSON in code blocks
        content = response.content.strip()
        if content.startswith("{"):
            assessment = json.loads(content)
        else:
            # Extract JSON from inside fenced block
            start = content.find("{")
            end = content.rfind("}") + 1
            assessment = json.loads(content[start:end])

        return {
            "verification_status": assessment.get("status", "needs_revision"),
            "rejection_reason": assessment.get("reason") if assessment.get("status") == "needs_revision" else None
        }
    except json.JSONDecodeError:
        return {"verification_status": "needs_revision", "rejection_reason": "Parse error. Please retry."}

from langgraph.graph import StateGraph, END
from state import ResearchState
from nodes.researcher import researcher_node
from nodes.writer import writer_node
from nodes.verifier import verifier_node

def should_continue(state: ResearchState) -> str:
    """
    Conditional edge: decides what happens after the Verifier runs.
    Returns the name of the next node.
    """
    status = state.get("verification_status")
    revision_count = state.get("revision_count", 0)
    max_revisions = state.get("max_revisions", 3)
    
    if status == "approved":
        print("[Router] ✓ Draft approved — proceeding to final output")
        return "format_output"
    
    if revision_count >= max_revisions:
        print(f"[Router] ⚠ Max revisions ({max_revisions}) reached — exiting with best draft")
        return "max_revisions_reached"
    
    print(f"[Router] ↩ Draft rejected — routing back to Writer (attempt {revision_count + 1})")
    return "writer"  # Route back for revision

def format_output_node(state: ResearchState) -> dict:
    """Formats the approved draft as the final output."""
    return {
        "final_document": state["draft"],
        "error": None
    }

def max_revisions_node(state: ResearchState) -> dict:
    """Handles the case where max revisions are reached."""
    return {
        "final_document": state.get("draft"),  # Use best available draft
        "error": f"Maximum revisions reached. Last rejection: {state.get('rejection_reason', 'Unknown')}"
    }

# Build the graph
def build_research_graph() -> StateGraph:
    workflow = StateGraph(ResearchState)
    
    # Add all nodes
    workflow.add_node("researcher", researcher_node)
    workflow.add_node("writer", writer_node)
    workflow.add_node("verifier", verifier_node)
    workflow.add_node("format_output", format_output_node)
    workflow.add_node("max_revisions_reached", max_revisions_node)
    
    # Set the entry point
    workflow.set_entry_point("researcher")
    
    # Linear edges (no branching)
    workflow.add_edge("researcher", "writer")
    workflow.add_edge("writer", "verifier")
    
    # Conditional edge after verification
    workflow.add_conditional_edges(
        "verifier",
        should_continue,
        {
            "format_output": "format_output",
            "writer": "writer",
            "max_revisions_reached": "max_revisions_reached"
        }
    )
    
    # Terminal edges
    workflow.add_edge("format_output", END)
    workflow.add_edge("max_revisions_reached", END)
    
    return workflow.compile()

# Example usage
if __name__ == "__main__":
    graph = build_research_graph()
    
    result = graph.invoke({
        "topic": "Multi-tenant vector database security best practices for B2B SaaS",
        "revision_count": 0,
        "max_revisions": 3,
        "research_findings": None,
        "draft": None,
        "verification_status": None,
        "rejection_reason": None,
        "final_document": None,
        "error": None
    })
    
    if result.get("error"):
        print(f"Completed with warning: {result['error']}")
    
    print("\n=== FINAL DOCUMENT ===")
    print(result["final_document"])