LangGraph Multi Agent Orchestration

%pip install -U langchain-community langgraph langchain-chroma langchain_aws pandas

import getpass
import os

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = getpass.getpass()
os.environ["LANGCHAIN_PROJECT"] = "lc-agent-sample"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"

from langchain_core.tools import tool
import random
from datetime import datetime, timedelta
from langgraph.prebuilt import ToolNode
import sqlite3


@tool
def search_flights(departure_city: str, arrival_city: str, date: str = None) -> str:
    """
    Use this tool to search for flights between two cities

    Args:
        departure_city (str): The city of departure
        arrival_city (str): The city of arrival
        date (str, optional): The date of the flight in YYYY-MM-DD format. If not provided, defaults to 7 days from now.

    Returns:
        str: A formatted string containing flight information including airline, departure time, arrival time, duration, and price for multiple flights.
    """
    departure_city = departure_city.capitalize()
    arrival_city = arrival_city.capitalize()

    if date is None:
        date = (datetime.now() + timedelta(days=7)).strftime("%Y-%m-%d")

    # Generate mock flight data
    num_flights = random.randint(2, 5)
    airlines = ["AirEurope", "SkyWings", "TransContinental", "EuroJet", "GlobalAir"]
    flights = []

    for _ in range(num_flights):
        airline = random.choice(airlines)
        duration = timedelta(minutes=2)
        price = random.randint(100, 400)
        departure_time = datetime.strptime(date, "%Y-%m-%d") + timedelta(
            hours=random.randint(0, 23), minutes=random.randint(0, 59)
        )
        arrival_time = departure_time + duration

        flights.append(
            {
                "airline": airline,
                "departure": departure_time.strftime("%H:%M"),
                "arrival": arrival_time.strftime("%H:%M"),
                "duration": str(duration),
                "price": price,
            }
        )

    # Format the results
    import json

    flight_data = {
        "departure_city": departure_city,
        "arrival_city": arrival_city,
        "date": date,
        "flights": []
    }
    for i, flight in enumerate(flights, 1):
        flight_info = {
            "flight_number": i,
            "airline": flight['airline'],
            "departure": flight['departure'],
            "arrival": flight['arrival'],
            "duration": str(flight['duration']),
            "price": flight['price']
        }
        flight_data["flights"].append(flight_info)

    return json.dumps(flight_data) + " FINISHED"

@tool
def retrieve_flight_booking(booking_id: str) -> str:
    """
    Retrieve a flight booking by ID

    Args:
        booking_id (str): The unique identifier of the booking to retrieve

    Returns:
        str: A string containing the booking information if found, or a message indicating no booking was found
    """
    conn = sqlite3.connect("data/travel_bookings.db")
    cursor = conn.cursor()

    # Execute the query to retrieve the booking
    cursor.execute("SELECT * FROM flight_bookings WHERE booking_id = ?", (booking_id,))
    booking = cursor.fetchone()

    # Close the connection
    conn.close()

    if booking:
        return f"Booking found: {booking}"
    else:
        return f"No booking found with ID: {booking_id} FINISHED"

@tool
def change_flight_booking(booking_id: str, new_date: str) -> str:
    """
    Change the date of a flight booking

    Args:
        booking_id (str): The unique identifier of the booking to be changed
        new_date (str): The new date for the booking

    Returns:
        str: A message indicating the result of the booking change operation
    """
    conn = sqlite3.connect("data/travel_bookings.db")
    cursor = conn.cursor()

    # Execute the query to update the booking date
    cursor.execute(
        "UPDATE flight_bookings SET departure_date = ? WHERE booking_id = ?",
        (new_date, booking_id),
    )
    conn.commit()

    # Check if the booking was updated
    if cursor.rowcount > 0:
        result = f"Booking updated with ID: {booking_id}, new date: {new_date} FINISHED"
    else:
        result = f"No booking found with ID: {booking_id} FINISHED"

    # Close the connection
    conn.close()

    return result 

@tool
def cancel_flight_booking(booking_id: str) -> str:
    """
    Cancel a flight booking. If the task complete, reply with "FINISHED"

    Args:
        booking_id (str): The unique identifier of the booking to be cancelled

    Returns:
        str: A message indicating the result of the booking cancellation operation

    """
    conn = sqlite3.connect("data/travel_bookings.db")
    cursor  = conn.cursor()

    cursor.execute("DELETE FROM flight_bookings WHERE booking_id = ?", (booking_id,))
    conn.commit()

    # Check if the booking was deleted
    if cursor.rowcount > 0:
        result = f"Booking canceled with ID: {booking_id} FINISHED"
    else:
        result = f"No booking found with ID: {booking_id} FINISHED"

    # Close the connection
    conn.close()

    return result

from langchain_aws import ChatBedrockConverse
from langchain_aws import ChatBedrock
import boto3

<h1>---- ⚠️ Update region for your AWS setup ⚠️ ----</h1>
bedrock_client = boto3.client("bedrock-runtime", region_name="us-east-1")

llm = ChatBedrockConverse(
    model="anthropic.claude-3-sonnet-20240229-v1:0",
    # model="anthropic.claude-3-5-sonnet-20240620-v1:0",
    temperature=0,
    max_tokens=None,
    client=bedrock_client,
    # other params...
)

import functools
import operator
from typing import Sequence, TypedDict

from langchain_core.messages import BaseMessage

from langgraph.graph import END, StateGraph, START
from langgraph.prebuilt import create_react_agent

from typing import Annotated
from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import MemorySaver

class State(TypedDict):
    messages: Annotated[list, add_messages]
    next: str

memory = MemorySaver()


flight_agent = create_react_agent(
    llm,
    tools=[
        search_flights,
        retrieve_flight_booking,
        change_flight_booking,
        cancel_flight_booking,
    ],
    state_modifier="""
    First gather all the information required to call a tool. 
    If you are not able to find the booking the do not try again and just reply with "FINISHED". 
    If tool has returned the results then reply with "FINISHED"
    If all tasks are complete, reply with "FINISHED"
    """,
    checkpointer=memory,
)

config = {"configurable": {"thread_id": "121"}}
flight_agent.invoke({"messages": [("user", "Can you give me booking details of booking number 10")]}, config)

import json

@tool
def suggest_hotels(city: str, checkin_date: str) -> dict:
    """
    Use this tool to search for hotels in these cities

    Args:
        city (str): The name of the city to search for hotels
        checkin_date (str): The check-in date in YYYY-MM-DD format

    Returns:
        dict: A dictionary containing:
            - hotels (list): List of hotel names in the specified city
            - checkin_date (str): The provided check-in date
            - checkout_date (str): A randomly generated checkout date
            - price (int): A randomly generated price for the stay
    """
    hotels = {
        "New York": ["Hotel A", "Hotel B", "Hotel C"],
        "Paris": ["Hotel D", "Hotel E", "Hotel F"],
        "Tokyo": ["Hotel G", "Hotel H", "Hotel I"],
    }

    # Generate random checkout date and price
    checkin = datetime.strptime(checkin_date, "%Y-%m-%d")
    checkout = checkin + timedelta(days=random.randint(1, 10))
    price = random.randint(100, 500)

    hotel_list = hotels.get(city, ["No hotels found"])
    hotel_data = {
        "hotels": hotel_list,
        "checkin_date": checkin_date,
        "checkout_date": checkout.strftime("%Y-%m-%d"),
        "price": price,
    }

    return json.dumps(hotel_data) + " FINISHED"

@tool
def retrieve_hotel_booking(booking_id: str) -> str:
    """
    Retrieve a hotel booking by ID

    Args:
        booking_id (str): The unique identifier of the hotel booking to retrieve

    Returns:
        str: A string containing the hotel booking information if found, or a message indicating no booking was found
    """
    conn = sqlite3.connect("data/travel_bookings.db")
    cursor = conn.cursor()
    cursor.execute(f"SELECT * FROM hotel_bookings WHERE booking_id='{booking_id}'")
    booking = cursor.fetchone()

    # Close the connection
    conn.close()

    if booking:
        return f"Booking found: {booking} FINISHED"
    else:
        return f"No booking found with ID: {booking_id} FINISHED"

from datetime import datetime

@tool
def change_hotel_booking(
    booking_id: int, new_checkin_date: str = None, new_checkout_date: str = None
) -> str:
    """
    Change the dates of a hotel booking in the database

    Args:
    booking_id (int): The unique identifier of the booking to be changed
    new_checkin_date (str, optional): The new check-in date in YYYY-MM-DD format
    new_checkout_date (str, optional): The new check-out date in YYYY-MM-DD format

    Returns:
    str: A message indicating the result of the booking change operation
    """

    conn = sqlite3.connect("data/travel_bookings.db")  # Replace with your actual database file
    cursor = conn.cursor()

    try:
        # First, fetch the current booking details
        cursor.execute(
            """
            SELECT * FROM hotel_bookings WHERE booking_id = ?
        """,
            (booking_id,),
        )

        booking = cursor.fetchone()

        if booking is None:
            return f"No hotel booking found with ID: {booking_id}"

        # Unpack the booking details
        (
            _,
            user_id,
            user_name,
            city,
            hotel_name,
            check_in_date,
            check_out_date,
            nights,
            price_per_night,
            total_price,
            num_guests,
            room_type,
        ) = booking

        # Update check-in and check-out dates if provided
        if new_checkin_date:
            check_in_date = new_checkin_date
        if new_checkout_date:
            check_out_date = new_checkout_date

        # Recalculate nights and total price
        checkin = datetime.strptime(check_in_date, "%Y-%m-%d")
        checkout = datetime.strptime(check_out_date, "%Y-%m-%d")
        nights = (checkout - checkin).days
        total_price = nights * price_per_night

        # Update the booking in the database
        cursor.execute(
            """
            UPDATE hotel_bookings
            SET check_in_date = ?, check_out_date = ?, nights = ?, total_price = ?
            WHERE booking_id = ?
        """,
            (check_in_date, check_out_date, nights, total_price, booking_id),
        )

        conn.commit()

        return f"Hotel booking updated: Booking ID {booking_id}, New check-in: {check_in_date}, New check-out: {check_out_date}, Nights: {nights}, Total Price: {total_price} FINISHED"

    except sqlite3.Error as e:
        conn.rollback()
        return f"An error occurred: {str(e)} FINISHED"

    finally:
        conn.close()

@tool
def cancel_hotel_booking(booking_id: str) -> str:
    """
    Cancel a hotel booking. If the task completes, reply with "FINISHED"

    Args:
        booking_id (str): The unique identifier of the booking to be cancelled

    Returns:
        str: A message indicating the result of the booking cancellation operation
    """
    conn = sqlite3.connect("data/travel_bookings.db")
    cursor  = conn.cursor()

    cursor.execute("DELETE FROM hotel_bookings WHERE booking_id = ?", (booking_id,))
    conn.commit()

    # Check if the booking was deleted
    if cursor.rowcount > 0:
        result = f"Booking canceled with ID: {booking_id} FINISHED"
    else:
        result = f"No booking found with ID: {booking_id} FINISHED"

    # Close the connection
    conn.close()

    return result

from langchain_core import __version__ as core_version
from packaging import version

core_version = version.parse(core_version)
if (core_version.major, core_version.minor) < (0, 3):
    from pydantic.v1 import BaseModel
else:
    from pydantic import BaseModel
from typing import Literal

class AskHuman(BaseModel):
    """Ask missing information from the user"""

    question: str

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable, RunnableConfig


class hotel_agent:
    def __init__(self, runnable: Runnable):
        self.runnable = runnable

    def __call__(self, state: State, config: RunnableConfig):
        while True:
            configuration = config.get("configurable", {})
            result = self.runnable.invoke(state)
            # If the LLM happens to return an empty response, we will re-prompt it
            # for an actual response.
            if not result.tool_calls and (
                not result.content
                or isinstance(result.content, list)
                and not result.content[0].get("text")
            ):
                # messages = state["messages"] + [("user", "Respond with a real output.")]
                messages = state["messages"]
                state = {**state, "messages": messages}
            else:
                break
        return {"messages": result}


primary_assistant_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful who manage hotel bookings"
            " If you dont have enough information use AskHuman tool to get additional information. "
            " City name is required to search hotels",
        ),
        ("placeholder", "{messages}"),
    ]
)
hotel_tools = [
    suggest_hotels,
    retrieve_hotel_booking,
    change_hotel_booking,
    cancel_hotel_booking,
]

runnable_with_tools = primary_assistant_prompt | llm.bind_tools(
    hotel_tools + [AskHuman]
)

<h1>We define a fake node to ask the human</h1>
def ask_human(state):
    pass

from langgraph.prebuilt import ToolNode, tools_condition

tool_node = ToolNode(hotel_tools+[AskHuman])

def should_continue(state):
    messages = state["messages"]
    last_message = messages[-1]
    # If there is no function call, then we finish
    if not last_message.tool_calls:
        return "end"
    elif last_message.tool_calls[0]["name"] == "AskHuman":
        return "ask_human"
    # Otherwise if there is, we continue
    else:
        return "continue"

from langgraph.graph import END, StateGraph, MessagesState
from IPython.display import Image, display

<h1>Define a new graph</h1>
workflow = StateGraph(MessagesState)

<h1>Define the three nodes we will cycle between</h1>
workflow.add_node("hotel_agent", hotel_agent(runnable_with_tools))
workflow.add_node("action", tool_node)
workflow.add_node("ask_human", ask_human)

workflow.add_edge(START, "hotel_agent")

<h1>We now add a conditional edge</h1>
workflow.add_conditional_edges(
    "hotel_agent",
    # Next, we pass in the function that will determine which node is called next.
    should_continue,
    {
        # If `tools`, then we call the tool node.
        "continue": "action",
        # We may ask the human
        "ask_human": "ask_human",
        # Otherwise we finish.
        "end": END,
    },
)


workflow.add_edge("action", "hotel_agent")

<h1>After we get back the human response, we go back to the agent</h1>
workflow.add_edge("ask_human", "hotel_agent")

<h1>Set up memory</h1>
from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()

app = workflow.compile(checkpointer=memory, interrupt_before=["ask_human"])

display(Image(app.get_graph().draw_mermaid_png()))

from langchain_core.messages import HumanMessage

config = {"configurable": {"thread_id": "128"}}
input_message = HumanMessage(content="I want to book a hotel for 8th October for 3 nights")
for event in app.stream({"messages": [input_message]}, config, stream_mode="values"):
    event["messages"][-1].pretty_print()

user_input = input("User: ")

tool_call_id = app.get_state(config).values["messages"][-1].tool_calls[0]["id"]

<h1>We now create the tool call with the id and the response we want</h1>
tool_message = [{"tool_call_id": tool_call_id, "type": "tool", "content": user_input}]

app.update_state(config, {"messages": tool_message}, as_node="ask_human")

app.get_state(config).next

for event in app.stream(None, config, stream_mode="values"):
    event["messages"][-1].pretty_print()

from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.graph.message import add_messages
from typing import Annotated, Sequence, List
from langchain_core.messages import HumanMessage

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core import __version__ as core_version
from packaging import version

core_version = version.parse(core_version)
if (core_version.major, core_version.minor) < (0, 3):
    from pydantic.v1 import BaseModel
else:
    from pydantic import BaseModel
from typing import Literal

members = ["flight_agent", "hotel_agent"]
options = ["FINISH"] + members

class routeResponse(BaseModel):
    """
    Return next agent name.
    """
    next: Literal[*options]


class AskHuman(BaseModel):
    """Ask missing information from the user"""

    question: str


prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """
        Given the conversation below who should act next?
        1. To search or cancel flight return 'flight_agent'
        2. To search for hotel or cancel hotel booking return 'hotel_agent'
        3. If you have the answer return 'FINISH'
        4. When member has finished the task, and you notice FINISHED in the message then don't repeat same member again
        5. Do not return next which is not related to user query. Example if user is asking about flight then do not call 'hotel_agent'
        Or should we FINISH? ONLY return one of these {options}. Do not explain the process.

        """,
        ),
        ("placeholder", "{messages}"),
    ]
).partial(options=str(options), members=", ".join(members))


class Supervisor:
    def __init__(self, runnable: Runnable):
        self.runnable = runnable

    def __call__(self, state: State, config: RunnableConfig):
        while True:
            configuration = config.get("configurable", {})
            result = self.runnable.invoke(state)
            # If the LLM happens to return an empty response, we will re-prompt it
            # for an actual response.
            if not result.tool_calls and (
                not result.content
                or isinstance(result.content, list)
                and not result.content[0].get("text")
            ):
                # messages = state["messages"] + [("user", "Respond with a real output.")]
                messages = state["messages"]
                state = {**state, "messages": messages}
            else:
                break
        return {"messages": result}


def supervisor_agent(state):
    supervisor_chain = prompt | llm.with_structured_output(
        routeResponse
    )
    result = supervisor_chain.invoke(state)
    print(result)
    output = {
        "next": result.next,
        "messages": [
            HumanMessage(
                content=f"Supervisor decided: {result.next}", name="supervisor"
            )
        ],
    }
    print(f"Supervisor output: {output}")
    return output

supervisor_agent({"messages": [("user", "I want to book a flight")]})

class State(TypedDict):
    messages: Annotated[list, add_messages]
    next: str

full_workflow = StateGraph(State)
full_workflow.add_node("supervisor", supervisor_agent)

full_workflow.add_edge(START, "supervisor")

from langchain_core.messages import AIMessage

def agent_node(state, agent, name):
    result = agent.invoke(state)
    return {
        "messages": [HumanMessage(content=result["messages"][-1].content, name=name)]
    }

flight_node = functools.partial(agent_node, agent=flight_agent, name="flight_agent")

full_workflow.add_node("flight_agent", flight_node)

full_workflow.add_node("hotel_agent", app)

def process_output(state):
    print(state)
    messages = state["messages"]
    for message in reversed(messages):
        if isinstance(message, AIMessage) and isinstance(message.content, str):
           return {
                "messages": [
                    HumanMessage(content=message.content, name="hotel_agent")
                ]
            }
    return None

full_workflow.add_node("process_output", process_output)

full_workflow.add_edge("flight_agent", "supervisor")
full_workflow.add_edge("hotel_agent", "process_output")
full_workflow.add_edge("process_output", "supervisor")
conditional_map = {k: k for k in members}
conditional_map["FINISH"] = END
full_workflow.add_conditional_edges("supervisor", lambda x: x["next"], conditional_map)

from IPython.display import Image, display

supervisor_agent_graph = full_workflow.compile(
    checkpointer=memory,
)

<h1>display subgraph using xray=1</h1>
display(Image(supervisor_agent_graph.get_graph(xray=1).draw_mermaid_png()))

config = {"configurable": {"thread_id": "133"}}
input_message = HumanMessage(
    content="I want to book a hotel for 8th October for 3 nights"
)
for event in supervisor_agent_graph.stream(
    {"messages": [input_message]}, config, stream_mode="values", subgraphs=True
):
    event[1]["messages"][-1].pretty_print()

state = supervisor_agent_graph.get_state(config, subgraphs=True)
state.tasks[0]

def extract_tool_id(pregel_task):
    # Navigate to the messages in the state
    messages = pregel_task.state.values.get("messages", [])

    # Find the last AIMessage
    for message in reversed(messages):
        if isinstance(message, AIMessage):
            # Check if the message has tool_calls
            tool_calls = getattr(message, "tool_calls", None)
            if tool_calls:
                # Return the id of the first tool call
                return tool_calls[0]["id"]

    # If no tool_id is found, return None
    return None


<h1>Assuming 'output' is the PregelTask object you provided</h1>
tool_id = extract_tool_id(state.tasks[0])
print(tool_id)  # This should print: tooluse_qWu6nEewS5OSqxszIgdfNA

user_input=input("User: ")

<h1>We now create the tool call with the id and the response we want</h1>
tool_message = [{"tool_call_id": tool_id, "type": "tool", "content": user_input}]


supervisor_agent_graph.update_state(
    state.tasks[0].state.config, {"messages": tool_message}, as_node="ask_human"
)

supervisor_agent_graph.get_state(state.tasks[0].state.config).next

for event in supervisor_agent_graph.stream(
   None, config, stream_mode="values", subgraphs=True
):
    event[1]["messages"][-1].pretty_print()