m2-l4.py

import asyncio

from langchain.agents import create_agent, AgentState
from langchain.messages import HumanMessage, ToolMessage
from langchain.tools import tool, ToolRuntime
from langchain_community.utilities import SQLDatabase
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.types import Command
from tavily import TavilyClient

QUESTION = HumanMessage(content="I'm from London and I'd like a wedding in Paris for 100 guests, jazz-genre")

tavily_client = TavilyClient()
client = MultiServerMCPClient(
    {
        "travel_server": {
            "transport": "streamable_http",
            "url": "https://mcp.kiwi.com",
        }
    }
)
db = SQLDatabase.from_uri("sqlite:///resources/Chinook.db")

@tool
def web_search(query: str) -> str:
    """
    Search the web for information.
    """
    return tavily_client.search(query)

@tool
def query_playlist_db(query: str) -> str:
    """
    Query the database for playlist information.
    """
    try:
        return db.run(query)
    except Exception as e:
        return f"Error querying the database: {e}"

class WeddingState(AgentState):
    origin: str
    destination: str
    guest_count: str
    genre: str

venue_agent = create_agent(
    model="gpt-5-nano",
    tools=[web_search],
    system_prompt="""
    You are a venue specialist. Search for venues in the desired location, and with the desired capacity.
    You are not allowed to ask any more follow up questions, you must find the best venue options based on the following criteria:
    - Price (lowest)
    - Capacity (exact match)
    - Reviews (highest)
    You may need to make multiple searches to iteratively find the best options.
    """
)

playlist_agent = create_agent(
    model="gpt-5-nano",
    tools=[query_playlist_db],
    system_prompt="""
    You are a playlist specialist. Query the sql database and curate the perfect playlist for a wedding given a genre.
    Once you have your playlist, calculate the total duration and cost of the playlist, each song has an associated price.
    If you run into errors when querying the database, try to fix them by making changes to the query.
    Do not come back empty handed, keep trying to query the db until you find a list of songs.
    You may need to make multiple queries to iteratively find the best options.
    """
)

@tool
def search_venues(runtime: ToolRuntime) -> str:
    """Venue agent chooses the best venue for the given location and capacity."""
    destination = runtime.state["destination"]
    capacity = runtime.state["guest_count"]
    query = f"Find wedding venues in {destination} for {capacity} guests"
    response = venue_agent.invoke({"messages": [HumanMessage(content=query)]})
    return response['messages'][-1].content

@tool
def suggest_playlist(runtime: ToolRuntime) -> str:
    """Playlist agent curates the perfect playlist for the given genre."""
    genre = runtime.state["genre"]
    query = f"Find {genre} tracks for wedding playlist"
    response = playlist_agent.invoke({"messages": [HumanMessage(content=query)]})
    return response['messages'][-1].content

@tool
def update_state(origin: str, destination: str, guest_count: str, genre: str, runtime: ToolRuntime) -> str:
    """Update the state when you know all of the values: origin, destination, guest_count, genre"""
    return Command(update={
        "origin": origin, 
        "destination": destination, 
        "guest_count": guest_count, 
        "genre": genre, 
        "messages": [ToolMessage("Successfully updated state", tool_call_id=runtime.tool_call_id)]}
        )

async def main():
    tools = await client.get_tools()
    travel_agent = create_agent(
        model="gpt-5-nano",
        tools=tools,
        system_prompt="""
        You are a travel agent. Search for flights to the desired destination wedding location.
        You are not allowed to ask any more follow up questions, you must find the best flight options based on the following criteria:
        - Price (lowest, economy class)
        - Duration (shortest)
        - Date (time of year which you believe is best for a wedding at this location)
        To make things easy, only look for one ticket, one way.
        You may need to make multiple searches to iteratively find the best options.
        You will be given no extra information, only the origin and destination. It is your job to think critically about the best options.
        Once you have found the best options, let the user know your shortlist of options.
        """
    )

    @tool
    async def search_flights(runtime: ToolRuntime) -> str:
        """Travel agent searches for flights to the desired destination wedding location."""
        origin = runtime.state["origin"]
        destination = runtime.state["destination"]
        response = await travel_agent.ainvoke({"messages": [HumanMessage(content=f"Find flights from {origin} to {destination}")]})
        return response['messages'][-1].content
    
    coordinator = create_agent(
         model="gpt-5-nano",
         tools=[search_flights, search_venues, suggest_playlist, update_state],
         state_schema=WeddingState,
         system_prompt="""
         You are a wedding coordinator. Delegate tasks to your specialists for flights, venues and playlists.
         First find all the information you need to update the state. Once that is done you can delegate the tasks.
         Once you have received their answers, coordinate the perfect wedding for me.
         """
    )
    response = await coordinator.ainvoke({"messages": [QUESTION]})
    print(response["messages"][-1].content)

if __name__ == "__main__":
    asyncio.run(main())