Building a LangGraph Multi-Agent System

mkdir multi-agent-system-using-langgraph
cd multi-agent-system-using-langgraph

# Create and activate virtual environment
uv venv
.venv\Scripts\activate

# Initialize project
uv init 

# Add dependencies
uv add langgraph langchain langchain_openai langchain_tavily toolbox-langchain

import random
from typing import TypedDict
from langgraph.graph import StateGraph, END

class ChainState(TypedDict):
    n: int
    number: int
    used_numbers: set[int]

def add(state: ChainState) -> ChainState:
  used_numbers = state.get("used_numbers", set())
  
  if len(used_numbers) >= 49:
        raise ValueError("All numbers from 1-49 have been used")
      
  available = set(range(1, 50)) - used_numbers          
  random_number = random.choice(list(available))
          
  new_used_numbers = used_numbers.copy()
  new_used_numbers.add(random_number)
    
  return {
        **state, 
        "n": state["n"] + 1, 
        "number": random_number,
        "used_numbers": new_used_numbers
    }

def dump(state: ChainState) -> ChainState:
    print("Current n:", state["n"], "number:", state["number"])
    print("Winning numbers:", sorted(state.get("used_numbers", set())))
    return state
    
def stop_condition(state: ChainState) -> bool:
    return state["n"] >= 6

graph = StateGraph(ChainState)
graph.add_node("add", add)
graph.add_node("dump", dump)
graph.add_edge("add", "dump")

graph.add_conditional_edges("dump", stop_condition, {
    True: END,
    False: "add",
})

graph.set_entry_point("add")
app = graph.compile()

result = app.invoke({"n": 0, "number": ""})

# Sample result
# Current n: 1 number: 43
# Winning numbers: [43]
# Current n: 2 number: 19
# Winning numbers: [19, 43]
# Current n: 3 number: 9
# Winning numbers: [9, 19, 43]
# Current n: 4 number: 36
# Winning numbers: [9, 19, 36, 43]
# Current n: 5 number: 38
# Winning numbers: [9, 19, 36, 38, 43]
# Current n: 6 number: 37
# Winning numbers: [9, 19, 36, 37, 38, 43]

from IPython.display import Image, display
display(Image(app.get_graph().draw_mermaid_png()))

CREATE TABLE hdb_combined_resale_flat_prices AS
SELECT month, town, flat_type, block, street_name, storey_range,
       floor_area_sqm, flat_model, lease_commence_date, resale_price
FROM raw_resale_flat_prices_from_jan_2017_onwards

UNION ALL SELECT * FROM raw_resale_flat_prices_from_jan_2015_to_dec_2016
UNION ALL SELECT * FROM raw_resale_flat_prices_from_mar_2012_to_dec_2014
UNION ALL SELECT * FROM raw_resale_flat_prices_2000_to_feb_2012
UNION ALL SELECT * FROM raw_resale_flat_prices_1990_1999;

sources:
  my-pg-source:
    kind: postgres
    host: postgres.local
    port: 5432
    database: postgres
    user: postgres
    password: postgres

tools:
  postgres-list-tables:
    kind: postgres-list-tables
    source: my-pg-source
    description: Retrieves schema information for all or specified tables.

  execute-sql-tool:
    kind: postgres-execute-sql
    source: my-pg-source
    description: Executes an arbitrary SQL statement.

  list-hdb-flats-by-town:
    kind: postgres-sql
    source: my-pg-source
    description: Returns the latest 5 resale transaction records for flats in a given town.
    parameters:
      - name: town
        type: string
        description: The town to search for.
    statement: >
      SELECT * 
      FROM public.hdb_combined_resale_flat_prices t
      WHERE t.town ILIKE '%' || $1 || '%'
      ORDER BY t.month DESC
      LIMIT 5;

  count-hdb-flats-by-town:
    kind: postgres-sql
    source: my-pg-source
    description: Counts the total number of resale transactions in the given town.
    parameters:
      - name: town
        type: string
        description: The town to count transactions in.
    statement: >
      SELECT COUNT(*) AS count
      FROM public.hdb_combined_resale_flat_prices t
      WHERE t.town ILIKE '%' || $1 || '%';

  percentile-price-by-town:
    kind: postgres-sql
    source: my-pg-source
    description: Returns a resale price percentile (median, p90, etc.) for a town.
    parameters:
      - name: town
        type: string
        description: The town to compute the percentile for.
      - name: percentile
        type: float
        description: The percentile to compute (0.5 = median).
        default: 0.5
    statement: >
      SELECT
        PERCENTILE_CONT($2) WITHIN GROUP (ORDER BY t.resale_price)
        AS percentile_price
      FROM public.hdb_combined_resale_flat_prices t
      WHERE t.town ILIKE '%' || $1 || '%';

  average-price-by-flat-type:
    kind: postgres-sql
    source: my-pg-source
    description: Returns the average resale price for a specific flat type, optionally filtered by town.
    parameters:
      - name: flat_type
        type: string
        description: The flat type (e.g., '3 ROOM', '4 ROOM').
      - name: town
        type: string
        description: Optional town filter (default = all towns).
        default: '%'
    statement: >
      SELECT AVG(t.resale_price) AS avg_price
      FROM public.hdb_combined_resale_flat_prices t
      WHERE t.flat_type = $1
      AND t.town ILIKE '%' || $2 || '%';

toolsets:
  my-toolset:
    - postgres-list-tables
    - execute-sql-tool
    - list-hdb-flats-by-town
    - count-hdb-flats-by-town
    - percentile-price-by-town
    - average-price-by-flat-type

$VERSION = "0.18.0"
Invoke-WebRequest -Uri "https://storage.googleapis.com/genai-toolbox/v$VERSION/windows/amd64/toolbox.exe" -OutFile "toolbox.exe"

.\toolbox

import os
import json
from typing import Annotated, Sequence, TypedDict

from langchain_core.messages import (
    HumanMessage, AIMessage, BaseMessage, ToolMessage
)
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from langchain_tavily import TavilySearch
from langchain.tools import tool

from toolbox_langchain import ToolboxClient

from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages

mcp_client = ToolboxClient("http://127.0.0.1:5000")

async def load_mcp():
    return await mcp_client.aload_toolset()

mcp_tools = await load_mcp()
print("Loaded MCP tools:", [t.name for t in mcp_tools])

tavily = TavilySearch(max_results=5)

@tool
def search_hdb_amenities(query: str) -> str:
    """
    Search for amenities near an HDB block or town using Tavily.

    Example queries:
    - "amenities near Block 123 Ang Mo Kio"
    - "schools near Bukit Panjang"
    - "supermarkets near Tampines"
    - "parks near Woodlands"

    The tool returns summarized search results.
    """
    formatted_query = f"amenities around {query}, Singapore HDB, nearby facilities, schools, supermarkets, malls, transport"
    results = tavily.invoke(formatted_query)
    return results

tools = [tavily, search_hdb_amenities] + mcp_tools
tools_by_name = {t.name: t for t in tools}

print("All tools loaded:", list(tools_by_name.keys()))

# Sample output
# All tools loaded: ['tavily_search', 'search_hdb_amenities', 'average-price-by-flat-type', 'count-hdb-flats-by-town', 'execute-sql-tool', 'list-hdb-flats-by-town', 'percentile-price-by-town', 'postgres-list-tables']

chat_prompt = ChatPromptTemplate.from_messages([
    ("system", """
You are an intelligent ReAct-style assistant with access to Postgres SQL (via MCP Toolbox)
and real-world web search tools (Tavily + amenities search).

Use SQL tools for:
- HDB transactions
- resale prices
- median / percentile
- flat type analysis
- town or block queries

Use search_hdb_amenities for:
- schools
- malls
- supermarkets
- parks
- hospitals
- amenities around blocks/towns

Use Tavily for general external knowledge.

Always:
1. Think step-by-step
2. Choose the correct tool
3. Call tools with correct arguments
4. Integrate tool output into your answer
"""),
    MessagesPlaceholder("scratch_pad")
])

MODEL = "moonshotai/kimi-k2-thinking"

model = ChatOpenAI(
    api_key=os.getenv("OPENROUTER_API_KEY"),
    base_url="https://openrouter.ai/api/v1",
    model=MODEL,
)

# Bind tools for ReAct
model_react = chat_prompt | model.bind_tools(tools)

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]

async def call_model(state: AgentState):
    response = await model_react.ainvoke({"scratch_pad": state["messages"]})
    return {"messages": [response]}

async def tool_node(state: AgentState):
    outputs = []
    last = state["messages"][-1]

    for tc in last.tool_calls:
        tool = tools_by_name[tc["name"]]

        # All toolbox_langchain tools support .ainvoke()
        result = await tool.ainvoke(tc["args"])

        outputs.append(
            ToolMessage(
                content=json.dumps(result),
                name=tc["name"],
                tool_call_id=tc["id"]
            )
        )

    return {"messages": outputs}

def should_continue(state: AgentState):
    last = state["messages"][-1]
    if last.tool_calls:
        return "continue"
    return "end"

workflow = StateGraph(AgentState)

workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)

workflow.add_edge("tools", "agent")

workflow.add_conditional_edges(
    "agent",
    should_continue,
    {
        "continue": "tools",
        "end": END,
    }
)

workflow.set_entry_point("agent")

graph = workflow.compile()

inputs = {
    "messages": [
        HumanMessage("Get me the past resale transactions for block 100 in Toa Payoh.")
    ]
}

async for step in graph.astream(inputs, stream_mode="values"):
    step["messages"][-1].pretty_print()

inputs = {
    "messages": [
        HumanMessage("What amenities are near Block 110 in Toa Payoh?")
    ]
}

async for step in graph.astream(inputs, stream_mode="values"):
    step["messages"][-1].pretty_print()

from IPython.display import Image, display
display(Image(graph.get_graph().draw_mermaid_png()))

import os
import json
from typing import Annotated, Sequence, TypedDict

from langchain_core.messages import (
    HumanMessage, AIMessage, BaseMessage, ToolMessage
)
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.tools import tool
from langchain_tavily import TavilySearch
from langchain_openai import ChatOpenAI

from toolbox_langchain import ToolboxClient
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages

mcp_client = ToolboxClient("http://127.0.0.1:5000")

async def load_mcp_tools():
    return await mcp_client.aload_toolset()

mcp_tools = await load_mcp_tools()
print("Loaded MCP tools:", [t.name for t in mcp_tools])

tavily = TavilySearch(max_results=5)

@tool
def search_hdb_amenities(query: str) -> str:
    """
    Search for amenities near an HDB block or town using Tavily.

    Example queries:
    - "amenities near Block 123 Ang Mo Kio"
    - "schools near Bukit Panjang"
    - "supermarkets near Tampines"
    - "parks near Woodlands"

    The tool returns summarized search results.
    """
    formatted_query = f"amenities around {query}, Singapore HDB, nearby facilities, schools, supermarkets, malls, transport"
    results = tavily.invoke(formatted_query)
    return results

all_tools = mcp_tools + [tavily, search_hdb_amenities]
tools_by_name = {t.name: t for t in all_tools}
print("All tools:", list(tools_by_name.keys()))

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]

sql_prompt = ChatPromptTemplate.from_messages([
    ("system", """
SQL Agent.
You answer only questions requiring HDB resale data or stats.

Use MCP SQL tools ONLY.
Always return structured, correct SQL tool calls.

"""),
    MessagesPlaceholder("scratch_pad")
])

amenities_prompt = ChatPromptTemplate.from_messages([
    ("system", """
Amenities Agent.
Use `search_hdb_amenities` first.
Use Tavily for deeper context.

You answer:
- amenities
- schools
- malls
- MRT
- supermarkets
near a block or town.
"""),
    MessagesPlaceholder("scratch_pad")
])

web_prompt = ChatPromptTemplate.from_messages([
    ("system", """
Web Research Agent.
Use Tavily search for general knowledge or non-HDB topics.
"""),
    MessagesPlaceholder("scratch_pad")
])

MODEL = "moonshotai/kimi-k2-thinking"

def build_model():
    return ChatOpenAI(
        api_key=os.getenv("OPENROUTER_API_KEY"),
        base_url="https://openrouter.ai/api/v1",
        model=MODEL,
    )

sql_agent = sql_prompt | build_model().bind_tools(mcp_tools)
amenities_agent = amenities_prompt | build_model().bind_tools([search_hdb_amenities, tavily])
web_agent = web_prompt | build_model().bind_tools([tavily])

supervisor_prompt = ChatPromptTemplate.from_messages([
    ("system", """
You are the Supervisor.

ROUTE the user query to one agent:
- SQL_AGENT for resale/transaction/statistics queries
- AMENITIES_AGENT for amenities/schools/MRT/malls
- WEB_AGENT for general queries

Respond with ONLY one of:
SQL_AGENT
AMENITIES_AGENT
WEB_AGENT
"""),
    MessagesPlaceholder("messages")
])

router_model = build_model()

async def router_node(state: AgentState):
    decision = await router_model.ainvoke(
        supervisor_prompt.format(messages=state["messages"])
    )

    route = decision.content.strip().upper()

    if route not in {"SQL_AGENT", "AMENITIES_AGENT", "WEB_AGENT"}:
        route = "WEB_AGENT"     # fallback so KeyError never happens

    return {"next_agent": route}

async def tool_node(state: AgentState):
    last = state["messages"][-1]
    outputs = []

    for tc in last.tool_calls:
        tool = tools_by_name[tc["name"]]
        result = await tool.ainvoke(tc["args"])
        
        outputs.append(ToolMessage(
            content=json.dumps(result),
            name=tc["name"],
            tool_call_id=tc["id"],
        ))

    return {"messages": outputs}

async def run_agent(agent, state):
    response = await agent.ainvoke({"scratch_pad": state["messages"]})
    return {"messages": [response]}

# Specialist agent wrappers
async def sql_agent_node(state):
    return await run_agent(sql_agent, state)

async def amenities_agent_node(state):
    return await run_agent(amenities_agent, state)

async def web_agent_node(state):
    return await run_agent(web_agent, state)

workflow = StateGraph(AgentState, async_mode=True)

# Register nodes
workflow.add_node("ROUTER", router_node)
workflow.add_node("SQL_AGENT", sql_agent_node)
workflow.add_node("AMENITIES_AGENT", amenities_agent_node)
workflow.add_node("WEB_AGENT", web_agent_node)
workflow.add_node("TOOLS", tool_node)

# Each agent → decide continue or end
def next_step(state: AgentState):
    last = state["messages"][-1]
    return "TOOLS" if last.tool_calls else END

workflow.add_conditional_edges("SQL_AGENT", next_step, {"TOOLS": "TOOLS", END: END})
workflow.add_conditional_edges("AMENITIES_AGENT", next_step, {"TOOLS": "TOOLS", END: END})
workflow.add_conditional_edges("WEB_AGENT", next_step, {"TOOLS": "TOOLS", END: END})


# Supervisor decides next hop
workflow.add_conditional_edges(
    "ROUTER",
    lambda out: out["next_agent"],
    {
        "SQL_AGENT": "SQL_AGENT",
        "AMENITIES_AGENT": "AMENITIES_AGENT",
        "WEB_AGENT": "WEB_AGENT",
    },
)

workflow.add_edge("TOOLS", "ROUTER")
workflow.set_entry_point("ROUTER")

graph = workflow.compile()

from IPython.display import Image, display
display(Image(graph.get_graph().draw_mermaid_png()))

inputs = {
    "messages": [
        HumanMessage("What amenities are near Block 110 in Toa Payoh?")
    ]
}

async for step in graph.astream(inputs, stream_mode="values"):
    step["messages"][-1].pretty_print()

async for step in graph.astream(
    {"messages": [HumanMessage("Show me the median resale price of 4-room flats in Bishan.")]},
    stream_mode="values"
):
    step["messages"][-1].pretty_print()

uv add langgraph-cli[inmem]

langgraph new app

cd app
python -m venv venv
.\venv\Scripts\Activate

pip install -e .

langgraph dev

What amenities are near Block 110 in Toa Payoh?