!!! note
    This notebook has been updated from the deprecated Swarm functionality to use the new GroupChat approach. The GroupChat system now accommodates any ConversableAgent and provides better orchestration through handoff conditions.

pip install --upgrade pydantic

import os

import autogen

llm_config = autogen.LLMConfig(
    config_list=[
        {
            "model": "gpt-5",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_type": "openai",
        },
        {
            "model": "gpt-5-mini",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_type": "openai",
        },
    ]
)
# Put the OpenAI API key into the environment
# os.environ["OPENAI_API_KEY"] = llm_config.config_list[0].api_key

from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType

# 3 Files (adjust path as necessary)
input_paths = [
    "../test/agentchat/contrib/graph_rag/trip_planner_data/attractions.jsonl",
    "../test/agentchat/contrib/graph_rag/trip_planner_data/cities.jsonl",
    "../test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.jsonl",
]
input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path) for input_path in input_paths]

from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation

# Attraction + Restaurant + City + Country Ontology
trip_data_ontology = Ontology()

trip_data_ontology.add_entity(
    Entity(
        label="Country",
        attributes=[
            Attribute(
                name="name",
                attr_type=AttributeType.STRING,
                required=True,
                unique=True,
            ),
        ],
    )
)
trip_data_ontology.add_entity(
    Entity(
        label="City",
        attributes=[
            Attribute(
                name="name",
                attr_type=AttributeType.STRING,
                required=True,
                unique=True,
            ),
            Attribute(
                name="weather",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
            Attribute(
                name="population",
                attr_type=AttributeType.NUMBER,
                required=False,
                unique=False,
            ),
        ],
    )
)
trip_data_ontology.add_entity(
    Entity(
        label="Restaurant",
        attributes=[
            Attribute(
                name="name",
                attr_type=AttributeType.STRING,
                required=True,
                unique=True,
            ),
            Attribute(
                name="description",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
            Attribute(
                name="rating",
                attr_type=AttributeType.NUMBER,
                required=False,
                unique=False,
            ),
            Attribute(
                name="food_type",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
        ],
    )
)
trip_data_ontology.add_entity(
    Entity(
        label="Attraction",
        attributes=[
            Attribute(
                name="name",
                attr_type=AttributeType.STRING,
                required=True,
                unique=True,
            ),
            Attribute(
                name="description",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
            Attribute(
                name="type",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
        ],
    )
)
trip_data_ontology.add_relation(
    Relation(
        label="IN_COUNTRY",
        source="City",
        target="Country",
    )
)
trip_data_ontology.add_relation(
    Relation(
        label="IN_CITY",
        source="Restaurant",
        target="City",
    )
)
trip_data_ontology.add_relation(
    Relation(
        label="IN_CITY",
        source="Attraction",
        target="City",
    )
)

from graphrag_sdk.models.openai import OpenAiGenerativeModel

from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import FalkorGraphQueryEngine
from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability

# Create FalkorGraphQueryEngine
query_engine = FalkorGraphQueryEngine(
    name="trip_data",
    # host="192.168.65.1",  # Change
    port=6379,  # if needed
    ontology=trip_data_ontology,
    model=OpenAiGenerativeModel("gpt-5-nano"),
)

# Ingest data and initialize the database
# query_engine.init_db(input_doc=input_documents)

# If you have already ingested and created the database, you can use this connect_db instead of init_db
query_engine.connect_db()

# IMPORTS
import json

import requests
from pydantic import BaseModel

from autogen import ConversableAgent, UserProxyAgent
from autogen.agentchat import initiate_group_chat
from autogen.agentchat.group import (
    AgentTarget,
    ContextVariables,
    OnCondition,
    ReplyResult,
    RevertToUserTarget,
    StringLLMCondition,
)
from autogen.agentchat.group.patterns import AutoPattern
from autogen.agentchat.group.targets.transition_target import AgentNameTarget

class Event(BaseModel):
    type: str  # Attraction, Restaurant, Travel
    location: str
    city: str
    description: str

class Day(BaseModel):
    events: list[Event]

class Itinerary(BaseModel):
    days: list[Day]

def _fetch_travel_time(origin: str, destination: str) -> dict:
    """Retrieves route information using Google Maps Directions API.
    API documentation at https://developers.google.com/maps/documentation/directions/get-directions
    """
    endpoint = "https://maps.googleapis.com/maps/api/directions/json"
    params = {
        "origin": origin,
        "destination": destination,
        "mode": "walking",  # driving (default), bicycling, transit
        "key": os.environ.get("GOOGLE_MAP_API_KEY"),
    }

    response = requests.get(endpoint, params=params)
    if response.status_code == 200:
        return response.json()
    else:
        return {"error": "Failed to retrieve the route information", "status_code": response.status_code}

def update_itinerary_with_travel_times(context_variables: ContextVariables) -> ReplyResult:
    """Update the complete itinerary with travel times between each event."""
    """
    Retrieves route information using Google Maps Directions API.
    API documentation at https://developers.google.com/maps/documentation/directions/get-directions
    """

    # Ensure that we have a structured itinerary, if not, back to the structured_output_agent to make it
    if context_variables.get("structured_itinerary") is None:
        print("DEBUG: No structured_itinerary found, redirecting to structured_output_agent")
        return ReplyResult(
            target=AgentNameTarget("structured_output_agent"),
            message="Structured itinerary not found, please create the structured output, structured_output_agent.",
        )
    elif "timed_itinerary" in context_variables:
        print("DEBUG: Timed itinerary already exists, task complete")
        # Return to user since work is done
        return ReplyResult(
            target=RevertToUserTarget(),
            context_variables=context_variables,
            message="Timed itinerary already done, inform the customer that their itinerary is ready!",
        )

    # Process the itinerary, converting it back to an object and working through each event to work out travel time and distance
    itinerary_object = Itinerary.model_validate(json.loads(context_variables["structured_itinerary"]))
    for day in itinerary_object.days:
        events = day.events
        new_events = []
        pre_event, cur_event = None, None
        event_count = len(events)
        index = 0
        while index < event_count:
            if index > 0:
                pre_event = events[index - 1]

            cur_event = events[index]
            if pre_event:
                origin = ", ".join([pre_event.location, pre_event.city])
                destination = ", ".join([cur_event.location, cur_event.city])
                maps_api_response = _fetch_travel_time(origin=origin, destination=destination)
                try:
                    leg = maps_api_response["routes"][0]["legs"][0]
                    travel_time_txt = f"{leg['duration']['text']}, ({leg['distance']['text']})"
                    new_events.append(
                        Event(
                            type="Travel",
                            location=f"walking from {pre_event.location} to {cur_event.location}",
                            city=cur_event.city,
                            description=travel_time_txt,
                        )
                    )
                except Exception:
                    print(f"Note: Unable to get travel time from {origin} to {destination}")
            new_events.append(cur_event)
            index += 1
        day.events = new_events

    context_variables["timed_itinerary"] = itinerary_object.model_dump()

    # Task complete - return to user with final context
    return ReplyResult(
        target=RevertToUserTarget(),
        context_variables=context_variables,
        message="Timed itinerary added to context with travel times. Your itinerary is ready!",
    )

trip_context = ContextVariables({
    "itinerary_confirmed": False,
    "itinerary": "",
    "structured_itinerary": None,
})

# Agent function placeholder - will be defined after agents are created
def mark_itinerary_as_complete(final_itinerary: str, context_variables: ContextVariables) -> ReplyResult:
    """Store and mark our itinerary as accepted by the customer."""

    context_variables["itinerary_confirmed"] = True
    context_variables["itinerary"] = final_itinerary

    # This will update the context variables and then transfer to the Structured Output agent
    return ReplyResult(
        target=AgentNameTarget("structured_output_agent"),
        context_variables=context_variables,
        message="Itinerary recorded and confirmed.",
    )

def create_structured_itinerary(context_variables: ContextVariables, structured_itinerary: str) -> ReplyResult:
    """Once a structured itinerary is created, store it and pass on to the Route Timing agent."""

    # Store the structured itinerary regardless of confirmation status
    # The confirmation should happen at the planner level, not here
    context_variables["structured_itinerary"] = structured_itinerary

    # This will update the context variables and then transfer to the Route Timing agent
    return ReplyResult(
        target=AgentNameTarget("route_timing_agent"),
        context_variables=context_variables,
        message="Structured itinerary stored.",
    )

# Planner agent, interacting with the customer and GraphRag agent, to create an itinerary
planner_agent = ConversableAgent(
    name="planner_agent",
    system_message="You are a trip planner agent. It is important to know where the customer is going, how many days, what they want to do."
    + "You will work with another agent, graphrag_agent, to get information about restaurant and attractions. "
    + "You are also working with the customer, so you must ask the customer what they want to do if you don't have LOCATION, NUMBER OF DAYS, MEALS, and ATTRACTIONS. "
    + "When you have the customer's requirements, work with graphrag_agent to get information for an itinerary."
    + "You are responsible for creating the itinerary and for each day in the itinerary you MUST HAVE events and EACH EVENT MUST HAVE a 'type' ('Restaurant' or 'Attraction'), 'location' (name of restaurant or attraction), 'city', and 'description'. "
    + "Finally, YOU MUST ask the customer if they are happy with the itinerary before marking the itinerary as complete.",
    functions=[mark_itinerary_as_complete],
    llm_config=llm_config,
)

# FalkorDB GraphRAG agent, utilising the FalkorDB to gather data for the Planner agent
graphrag_agent = ConversableAgent(
    name="graphrag_agent",
    system_message="Return a list of restaurants and/or attractions. List them separately and provide ALL the options in the location. Do not provide travel advice.",
)

# Adding the FalkorDB capability to the agent
graph_rag_capability = FalkorGraphRagCapability(query_engine)
graph_rag_capability.add_to_agent(graphrag_agent)

# Structured Output agent, formatting the itinerary into a structured format through the response_format on the LLM Configuration
structured_llm_config = llm_config = autogen.LLMConfig(
    config_list=[
        {
            "model": "gpt-5",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_type": "openai",
        },
        {
            "model": "gpt-5-mini",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_type": "openai",
        },
    ]
)

for config in structured_llm_config.config_list:
    config.response_format = Itinerary

structured_output_agent = ConversableAgent(
    name="structured_output_agent",
    system_message="You are a data formatting agent, format the provided itinerary in the context below into the provided format. IMPORTANT: You must call the create_structured_itinerary function to save the formatted itinerary.",
    llm_config=structured_llm_config,
    functions=[create_structured_itinerary],
)

# Route Timing agent, adding estimated travel times to the itinerary by utilising the Google Maps Platform
route_timing_agent = ConversableAgent(
    name="route_timing_agent",
    system_message="You are a route timing agent. Your job is to call the update_itinerary_with_travel_times tool to add travel times to the itinerary. The tool will handle everything including ending the conversation when done.",
    llm_config=llm_config,
    functions=[update_itinerary_with_travel_times],
)

# Our customer will be a human in the loop
customer = UserProxyAgent(name="customer")

# Set up handoffs for each agent using the proper handoffs methods
# NOTE: For handoffs, we use AgentTarget with actual agent objects
planner_agent.handoffs.add_llm_conditions([
    OnCondition(
        target=AgentTarget(graphrag_agent),
        condition=StringLLMCondition("Need information on the restaurants and attractions for a location"),
    ),
    OnCondition(
        target=AgentTarget(structured_output_agent),
        condition=StringLLMCondition("Itinerary is confirmed by the customer"),
    ),
])

# Back to the Planner when information has been retrieved
graphrag_agent.handoffs.set_after_work(AgentTarget(planner_agent))

# Once we have formatted our itinerary, we can hand off to the route timing agent to add in the travel timings
structured_output_agent.handoffs.set_after_work(AgentTarget(route_timing_agent))

# Route timing agent will revert to user when done (handled by the tool function)
# Note: No explicit handoff needed since the tool handles returning to user

# Create AutoPattern for dynamic agent selection
pattern = AutoPattern(
    initial_agent=planner_agent,
    agents=[planner_agent, graphrag_agent, structured_output_agent, route_timing_agent],
    user_agent=customer,
    context_variables=trip_context,
    group_manager_args={"llm_config": llm_config},
    # group_after_work=RevertToUserTarget(),
)

# Start the conversation using GroupChat with AutoPattern
chat_result, final_context, last_agent = initiate_group_chat(
    pattern=pattern,
    messages="I want to go to Rome for a couple of days. Can you help me plan my trip?",
    max_rounds=100,
)

def print_itinerary(itinerary_data):
    header = "█             █\\n █           █ \\n  █  █████  █  \\n   ██     ██   \\n  █         █  \\n █  ███████  █ \\n █ ██ ███ ██ █ \\n   █████████   \\n\\n ██   ███ ███  \\n█  █ █       █ \\n████ █ ██  ██  \\n█  █ █  █ █    \\n█  █  ██  ████ \\n"
    width = 80
    icons = {"Travel": "🚶", "Restaurant": "🍽️", "Attraction": "🏛️"}

    for line in header.split("\\n"):
        print(line.center(width))
    print(f"Itinerary for {itinerary_data['days'][0]['events'][0]['city']}".center(width))
    print("=" * width)

    for day_num, day in enumerate(itinerary_data["days"], 1):
        print(f"\\nDay {day_num}".center(width))
        print("-" * width)

        for event in day["events"]:
            event_type = event["type"]
            print(f"\\n  {icons[event_type]} {event['location']}")
            if event_type != "Travel":
                words = event["description"].split()
                line = "    "
                for word in words:
                    if len(line) + len(word) + 1 <= 76:
                        line += word + " "
                    else:
                        print(line)
                        line = "    " + word + " "
                if line.strip():
                    print(line)
            else:
                print(f"    {event['description']}")
        print("\\n" + "-" * width)

# First, try to get the itinerary from context variables
print("Checking context variables for itinerary data...")

# Check multiple potential sources for the itinerary
itinerary_data = None
itinerary_found = False

# 1. Check final_context
if final_context.get("timed_itinerary"):
    print("Found timed_itinerary in final_context!")
    itinerary_data = final_context["timed_itinerary"]
    itinerary_found = True
elif final_context.get("structured_itinerary"):
    print("Found structured_itinerary in final_context!")
    itinerary_data = json.loads(final_context["structured_itinerary"])
    itinerary_found = True

# 2. Check pattern context variables if final_context didn't have it
if not itinerary_found and hasattr(pattern, "context_variables"):
    pattern_context = pattern.context_variables
    if pattern_context.get("timed_itinerary"):
        print("Found timed_itinerary in pattern context!")
        itinerary_data = pattern_context["timed_itinerary"]
        itinerary_found = True
    elif pattern_context.get("structured_itinerary"):
        print("Found structured_itinerary in pattern context!")
        itinerary_data = json.loads(pattern_context["structured_itinerary"])
        itinerary_found = True

# 3. If context variables failed, extract from chat messages as fallback
if not itinerary_found:
    print("Context variables empty, extracting itinerary from chat messages...")

    # Look through the chat result messages for structured output
    if hasattr(chat_result, "chat_history"):
        messages = chat_result.chat_history
    elif hasattr(chat_result, "messages"):
        messages = chat_result.messages
    else:
        messages = []

    print(f"Checking {len(messages)} messages for itinerary data...")

    for i, message in enumerate(messages):
        try:
            # Check if this message contains structured itinerary JSON
            content = message.get("content", "") if isinstance(message, dict) else str(message)

            # Look for JSON-like content that matches our Itinerary structure
            if "days" in content and "events" in content and "type" in content:
                print(f"Found potential itinerary in message {i}")

                # Try to extract JSON from the content
                import re

                json_match = re.search(r'\\{.*"days".*\\}', content, re.DOTALL)
                if json_match:
                    try:
                        itinerary_json = json_match.group(0)
                        itinerary_data = json.loads(itinerary_json)
                        print("Successfully parsed itinerary JSON from messages!")
                        itinerary_found = True
                        break
                    except json.JSONDecodeError:
                        continue
        except Exception:
            continue

# Display the itinerary if found
if itinerary_found and itinerary_data:
    print("\\n" + "=" * 80)
    print("ITINERARY FOUND! Printing...")
    print("=" * 80)
    print_itinerary(itinerary_data)
else:
    print("No valid itinerary found anywhere.")
    print("\\nDEBUG INFO:")
    print("Final context:", final_context.to_dict() if hasattr(final_context, "to_dict") else final_context)
    print(
        "Pattern context:",
        pattern.context_variables.to_dict()
        if hasattr(pattern, "context_variables") and hasattr(pattern.context_variables, "to_dict")
        else "No pattern context",
    )

    # Show available keys in contexts
    if hasattr(final_context, "keys") or hasattr(final_context, "to_dict"):
        available_keys = (
            list(final_context.keys())
            if hasattr(final_context, "keys")
            else list(final_context.to_dict().keys())
            if hasattr(final_context, "to_dict")
            else "Unknown"
        )
        print(f"Available keys in final_context: {available_keys}")

    if hasattr(pattern, "context_variables"):
        pattern_keys = (
            list(pattern.context_variables.keys())
            if hasattr(pattern.context_variables, "keys")
            else list(pattern.context_variables.to_dict().keys())
            if hasattr(pattern.context_variables, "to_dict")
            else "Unknown"
        )
        print(f"Available keys in pattern context: {pattern_keys}")