notebook/agentchat_swarm_graphrag_telemetry_trip_planner.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using a local Telemetry server to monitor a GraphRAG agent\n",
    "\n",
    "In this notebook, we're building a trip planning swarm which has an objective to create an itinerary together with a customer. The end result will be an itinerary that has route times and distances calculated between activities.\n",
    "\n",
    "The following diagram outlines the key components of the Swarm, with highlights being:\n",
    "\n",
    "- **Arize Phoenix to provide transparency using the OpenTelemetry standard**\n",
    "- FalkorDB agent using a GraphRAG database of restaurants and attractions\n",
    "- Structured Output agent that will enforce a strict format for the accepted itinerary\n",
    "- Routing agent that utilises the Google Maps API to calculate distances between activities\n",
    "- Swarm orchestration utilising context variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````mdx-code-block\n",
    ":::note\n",
    "This notebook has been updated as swarms can now accommodate any ConversableAgent.\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize Python environment\n",
    "\n",
    "- FalkorDB's GraphRAG-SDK is a dependency for this notebook\n",
    "- Please ensure you have Pydantic version 2+ installed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qU pydantic ag2[graph-rag-falkor-db] ipykernel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Docker Containers\n",
    "\n",
    "**Note:** This may require a Docker Compose file to get it to work reliably."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FalkorDB\n",
    "\n",
    "- UI endpoint:  http://localhost:3000/graph\n",
    "- sample query:  `match path = ()-[]-() return path`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if you run the FalkorDB image with the rm flag it is removed after the container is stopped\n",
    "# for more information refer to:  https://docs.falkordb.com/\n",
    "\n",
    "# !docker run -p 6379:6379 -p 3000:3000 -it --rm falkordb/falkordb:latest\n",
    "\n",
    "import subprocess\n",
    "\n",
    "# Run the Docker container without interactive mode\n",
    "subprocess.Popen([\"docker\", \"run\", \"-p\", \"6379:6379\", \"-p\", \"3000:3000\", \"--rm\", \"falkordb/falkordb:latest\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Arize Phoenix\n",
    "\n",
    "- UI endpoint:  http://localhost:6006"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for more information refer to https://docs.arize.com/phoenix/tracing/integrations-tracing/autogen-support#docker\n",
    "# !docker run -p 6006:6006 -p 4317:4317 arizephoenix/phoenix:latest\n",
    "\n",
    "import subprocess\n",
    "\n",
    "# Run the Docker container without interactive mode\n",
    "subprocess.Popen([\"docker\", \"run\", \"-p\", \"6006:6006\", \"-p\", \"4317:4317\", \"--rm\", \"arizephoenix/phoenix:latest\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Arize Phoenix:  setup and configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q arize-phoenix-otel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phoenix.otel import register\n",
    "\n",
    "# defaults to endpoint=\"http://localhost:4317\"\n",
    "tracer_provider = register(\n",
    "    project_name=\"ag2-swarm-graphrag\",  # Default is 'default'\n",
    "    endpoint=\"http://localhost:4317\",  # Sends traces using gRPC\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install python telemetry library requirements\n",
    "!pip install -q openinference-instrumentation-openai 'httpx<0.28'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openinference.instrumentation.openai import OpenAIInstrumentor\n",
    "\n",
    "OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Google Maps API Key\n",
    "\n",
    "To use Google's API to calculate travel times, you will need to have enabled the `Directions API` in your Google Maps Platform. You can get an API key and free quota, see [here](https://developers.google.com/maps/documentation/directions/overview) and [here](https://developers.google.com/maps/get-started) for more details.\n",
    "\n",
    "Once you have your API key, set the environment variable `GOOGLE_MAP_API_KEY` to this value.\n",
    "\n",
    "NOTE:  If the route planning step is failing, it is likely an environment variable issue which can occur in Jupyter notebooks.  The code in the `update_itinerary_with_travel_times` and  `_fetch_travel_time` functions below could be enhanced to provide better visibility if these API calls fail.  The following code cell can assist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"GOOGLE_MAP_API_KEY\"] = os.getenv(\"GOOGLE_MAP_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set Configuration and OpenAI API Key\n",
    "\n",
    "**Create a OAI_CONFIG_LIST file** in the AG2 project `notebook` directory based on the OAI_CONFIG_LIST_sample file from the root directory.\n",
    "\n",
    "By default, FalkorDB uses OpenAI LLMs and that requires an OpenAI key in your environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "You can utilise an OAI_CONFIG_LIST file and extract the OpenAI API key and put it in the environment, as will be shown in the following cell.\n",
    "\n",
    "Alternatively, you can load the environment variable yourself.\n",
    "\n",
    "````{=mdx}\n",
    ":::tip\n",
    "Learn more about configuring LLMs for agents [here](https://docs.ag2.ai/docs/user-guide/basic-concepts/llm-configuration/llm-configuration)llm-configuration).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import autogen\n",
    "\n",
    "llm_config = autogen.LLMConfig.from_json(path=\"OAI_CONFIG_LIST\", timeout=120).where(tags=[\"gpt-4o\"])\n",
    "\n",
    "# Put the OpenAI API key into the environment using the config_list or env variable\n",
    "# os.environ[\"OPENAI_API_KEY\"] = llm_config.config_list[0].api_key\n",
    "os.environ[\"OPENAI_API_KEY\"] = os.getenv(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare the FalkorDB GraphRAG database\n",
    "\n",
    "Using 3 sample JSON data files from our GitHub repository, we will create a specific ontology for our GraphRAG database and then populate it.\n",
    "\n",
    "Creating a specific ontology that matches with the types of queries makes for a more optimal database and is more cost efficient when populating the knowledge graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType\n",
    "\n",
    "# 3 Files (adjust path as necessary)\n",
    "input_paths = [\n",
    "    \"../test/agentchat/contrib/graph_rag/trip_planner_data/attractions.jsonl\",\n",
    "    \"../test/agentchat/contrib/graph_rag/trip_planner_data/cities.jsonl\",\n",
    "    \"../test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.jsonl\",\n",
    "]\n",
    "input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path) for input_path in input_paths]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Ontology\n",
    "\n",
    "Entities: Country, City, Attraction, Restaurant\n",
    "\n",
    "Relationships: City in Country, Attraction in City, Restaurant in City"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation\n",
    "\n",
    "# Attraction + Restaurant + City + Country Ontology\n",
    "trip_data_ontology = Ontology()\n",
    "\n",
    "trip_data_ontology.add_entity(\n",
    "    Entity(\n",
    "        label=\"Country\",\n",
    "        attributes=[\n",
    "            Attribute(\n",
    "                name=\"name\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=True,\n",
    "                unique=True,\n",
    "            ),\n",
    "        ],\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_entity(\n",
    "    Entity(\n",
    "        label=\"City\",\n",
    "        attributes=[\n",
    "            Attribute(\n",
    "                name=\"name\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=True,\n",
    "                unique=True,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"weather\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"population\",\n",
    "                attr_type=AttributeType.NUMBER,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "        ],\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_entity(\n",
    "    Entity(\n",
    "        label=\"Restaurant\",\n",
    "        attributes=[\n",
    "            Attribute(\n",
    "                name=\"name\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=True,\n",
    "                unique=True,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"description\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"rating\",\n",
    "                attr_type=AttributeType.NUMBER,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"food_type\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "        ],\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_entity(\n",
    "    Entity(\n",
    "        label=\"Attraction\",\n",
    "        attributes=[\n",
    "            Attribute(\n",
    "                name=\"name\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=True,\n",
    "                unique=True,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"description\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "            Attribute(\n",
    "                name=\"type\",\n",
    "                attr_type=AttributeType.STRING,\n",
    "                required=False,\n",
    "                unique=False,\n",
    "            ),\n",
    "        ],\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_relation(\n",
    "    Relation(\n",
    "        label=\"IN_COUNTRY\",\n",
    "        source=\"City\",\n",
    "        target=\"Country\",\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_relation(\n",
    "    Relation(\n",
    "        label=\"IN_CITY\",\n",
    "        source=\"Restaurant\",\n",
    "        target=\"City\",\n",
    "    )\n",
    ")\n",
    "trip_data_ontology.add_relation(\n",
    "    Relation(\n",
    "        label=\"IN_CITY\",\n",
    "        source=\"Attraction\",\n",
    "        target=\"City\",\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize FalkorDB and Query Engine\n",
    "\n",
    "Remember: Change your host, port, and preferred OpenAI model if needed (gpt-4o-mini and better is recommended)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphrag_sdk.models.openai import OpenAiGenerativeModel\n",
    "\n",
    "from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import FalkorGraphQueryEngine\n",
    "from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability\n",
    "\n",
    "# Create FalkorGraphQueryEngine\n",
    "query_engine = FalkorGraphQueryEngine(\n",
    "    name=\"trip_data\",\n",
    "    host=\"localhost\",  # change to a specific IP address if you run into issues connecting to your local instance\n",
    "    port=6379,  # if needed\n",
    "    ontology=trip_data_ontology,\n",
    "    model=OpenAiGenerativeModel(\"gpt-4o\"),\n",
    ")\n",
    "\n",
    "# Ingest data and initialize the database\n",
    "query_engine.init_db(input_doc=input_documents)\n",
    "\n",
    "# If you have already ingested and created the database, you can use this connect_db instead of init_db\n",
    "# query_engine.connect_db()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IMPORTS\n",
    "import json\n",
    "import os\n",
    "from typing import Any, Dict\n",
    "\n",
    "import requests\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from autogen import (\n",
    "    AfterWork,\n",
    "    AfterWorkOption,\n",
    "    ConversableAgent,\n",
    "    OnCondition,\n",
    "    SwarmResult,\n",
    "    UserProxyAgent,\n",
    "    initiate_swarm_chat,\n",
    "    register_hand_off,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pydantic model for Structured Output\n",
    "\n",
    "Utilising OpenAI's [Structured Outputs](https://platform.openai.com/docs/guides/structured-outputs), our Structured Output agent's responses will be constrained to this Pydantic model.\n",
    "\n",
    "The itinerary is structured as:\n",
    "Itinerary has Day(s) has Event(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Event(BaseModel):\n",
    "    type: str  # Attraction, Restaurant, Travel\n",
    "    location: str\n",
    "    city: str\n",
    "    description: str\n",
    "\n",
    "\n",
    "class Day(BaseModel):\n",
    "    events: list[Event]\n",
    "\n",
    "\n",
    "class Itinerary(BaseModel):\n",
    "    days: list[Day]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Google Maps Platform\n",
    "\n",
    "The functions necessary to query the Directions API to get travel times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _fetch_travel_time(origin: str, destination: str) -> dict:\n",
    "    \"\"\"Retrieves route information using Google Maps Directions API.\n",
    "    API documentation at https://developers.google.com/maps/documentation/directions/get-directions\n",
    "    \"\"\"\n",
    "    endpoint = \"https://maps.googleapis.com/maps/api/directions/json\"\n",
    "    params = {\n",
    "        \"origin\": origin,\n",
    "        \"destination\": destination,\n",
    "        \"mode\": \"walking\",  # driving (default), bicycling, transit\n",
    "        \"key\": os.environ.get(\"GOOGLE_MAP_API_KEY\"),\n",
    "    }\n",
    "\n",
    "    response = requests.get(endpoint, params=params)\n",
    "    if response.status_code == 200:\n",
    "        return response.json()\n",
    "    else:\n",
    "        return {\"error\": \"Failed to retrieve the route information\", \"status_code\": response.status_code}\n",
    "\n",
    "\n",
    "def update_itinerary_with_travel_times(context_variables: dict) -> SwarmResult:\n",
    "    \"\"\"Update the complete itinerary with travel times between each event.\"\"\"\n",
    "    \"\"\"\n",
    "    Retrieves route information using Google Maps Directions API.\n",
    "    API documentation at https://developers.google.com/maps/documentation/directions/get-directions\n",
    "    \"\"\"\n",
    "\n",
    "    # Ensure that we have a structured itinerary, if not, back to the structured_output_agent to make it\n",
    "    if context_variables.get(\"structured_itinerary\") is None:\n",
    "        return SwarmResult(\n",
    "            agent=\"structured_output_agent\",\n",
    "            values=\"Structured itinerary not found, please create the structured output, structured_output_agent.\",\n",
    "        )\n",
    "    elif \"timed_itinerary\" in context_variables:\n",
    "        return SwarmResult(values=\"Timed itinerary already done, inform the customer that their itinerary is ready!\")\n",
    "\n",
    "    # Process the itinerary, converting it back to an object and working through each event to work out travel time and distance\n",
    "    itinerary_object = Itinerary.model_validate(json.loads(context_variables[\"structured_itinerary\"]))\n",
    "    for day in itinerary_object.days:\n",
    "        events = day.events\n",
    "        new_events = []\n",
    "        pre_event, cur_event = None, None\n",
    "        event_count = len(events)\n",
    "        index = 0\n",
    "        while index < event_count:\n",
    "            if index > 0:\n",
    "                pre_event = events[index - 1]\n",
    "\n",
    "            cur_event = events[index]\n",
    "            if pre_event:\n",
    "                origin = \", \".join([pre_event.location, pre_event.city])\n",
    "                destination = \", \".join([cur_event.location, cur_event.city])\n",
    "                maps_api_response = _fetch_travel_time(origin=origin, destination=destination)\n",
    "                try:\n",
    "                    leg = maps_api_response[\"routes\"][0][\"legs\"][0]\n",
    "                    travel_time_txt = f\"{leg['duration']['text']}, ({leg['distance']['text']})\"\n",
    "                    new_events.append(\n",
    "                        Event(\n",
    "                            type=\"Travel\",\n",
    "                            location=f\"walking from {pre_event.location} to {cur_event.location}\",\n",
    "                            city=cur_event.city,\n",
    "                            description=travel_time_txt,\n",
    "                        )\n",
    "                    )\n",
    "                except Exception:\n",
    "                    print(f\"Note: Unable to get travel time from {origin} to {destination}\")\n",
    "            new_events.append(cur_event)\n",
    "            index += 1\n",
    "        day.events = new_events\n",
    "\n",
    "    context_variables[\"timed_itinerary\"] = itinerary_object.model_dump()\n",
    "\n",
    "    return SwarmResult(context_variables=context_variables, values=\"Timed itinerary added to context with travel times\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Swarm\n",
    "\n",
    "### Context Variables\n",
    "Our swarm agents will have access to a couple of context variables in relation to the itinerary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trip_context = {\n",
    "    \"itinerary_confirmed\": False,\n",
    "    \"itinerary\": \"\",\n",
    "    \"structured_itinerary\": None,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agent Functions\n",
    "\n",
    "We have two functions/tools for our agents.\n",
    "\n",
    "One for our Planner agent to mark an itinerary as confirmed by the customer and to store the final text itinerary. This will then transfer to our Structured Output agent.\n",
    "\n",
    "Another for the Structured Output Agent to save the structured itinerary and transfer to the Route Timing agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mark_itinerary_as_complete(final_itinerary: str, context_variables: Dict[str, Any]) -> SwarmResult:\n",
    "    \"\"\"Store and mark our itinerary as accepted by the customer.\"\"\"\n",
    "    context_variables[\"itinerary_confirmed\"] = True\n",
    "    context_variables[\"itinerary\"] = final_itinerary\n",
    "\n",
    "    # This will update the context variables and then transfer to the Structured Output agent\n",
    "    return SwarmResult(\n",
    "        agent=\"structured_output_agent\", context_variables=context_variables, values=\"Itinerary recorded and confirmed.\"\n",
    "    )\n",
    "\n",
    "\n",
    "def create_structured_itinerary(context_variables: Dict[str, Any], structured_itinerary: str) -> SwarmResult:\n",
    "    \"\"\"Once a structured itinerary is created, store it and pass on to the Route Timing agent.\"\"\"\n",
    "    # Ensure the itinerary is confirmed, if not, back to the Planner agent to confirm it with the customer\n",
    "    if not context_variables[\"itinerary_confirmed\"]:\n",
    "        return SwarmResult(\n",
    "            agent=\"planner_agent\",\n",
    "            values=\"Itinerary not confirmed, please confirm the itinerary with the customer first.\",\n",
    "        )\n",
    "\n",
    "    context_variables[\"structured_itinerary\"] = structured_itinerary\n",
    "\n",
    "    # This will update the context variables and then transfer to the Route Timing agent\n",
    "    return SwarmResult(\n",
    "        agent=\"route_timing_agent\", context_variables=context_variables, values=\"Structured itinerary stored.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agents\n",
    "\n",
    "Our Swarm agents and a UserProxyAgent (human) which the swarm will interact with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Planner agent, interacting with the customer and GraphRag agent, to create an itinerary\n",
    "planner_agent = ConversableAgent(\n",
    "    name=\"planner_agent\",\n",
    "    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.\"\n",
    "    + \"You will work with another agent, graphrag_agent, to get information about restaurant and attractions. \"\n",
    "    + \"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. \"\n",
    "    + \"When you have the customer's requirements, work with graphrag_agent to get information for an itinerary.\"\n",
    "    + \"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'. \"\n",
    "    + \"Finally, YOU MUST ask the customer if they are happy with the itinerary before marking the itinerary as complete.\",\n",
    "    functions=[mark_itinerary_as_complete],\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "# FalkorDB GraphRAG agent, utilising the FalkorDB to gather data for the Planner agent\n",
    "graphrag_agent = ConversableAgent(\n",
    "    name=\"graphrag_agent\",\n",
    "    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.\",\n",
    ")\n",
    "\n",
    "# Adding the FalkorDB capability to the agent\n",
    "graph_rag_capability = FalkorGraphRagCapability(query_engine)\n",
    "graph_rag_capability.add_to_agent(graphrag_agent)\n",
    "\n",
    "# Structured Output agent, formatting the itinerary into a structured format through the response_format on the LLM Configuration\n",
    "structured_llm_config = autogen.LLMConfig.from_json(path=\"OAI_CONFIG_LIST\", timeout=120).where(tags=[\"gpt-4o\"])\n",
    "\n",
    "for config in structured_llm_config.config_list:\n",
    "    config.response_format = Itinerary\n",
    "\n",
    "structured_output_agent = ConversableAgent(\n",
    "    name=\"structured_output_agent\",\n",
    "    system_message=\"You are a data formatting agent, format the provided itinerary in the context below into the provided format.\",\n",
    "    llm_config=structured_llm_config,\n",
    "    functions=[create_structured_itinerary],\n",
    ")\n",
    "\n",
    "# Route Timing agent, adding estimated travel times to the itinerary by utilising the Google Maps Platform\n",
    "route_timing_agent = ConversableAgent(\n",
    "    name=\"route_timing_agent\",\n",
    "    system_message=\"You are a route timing agent. YOU MUST call the update_itinerary_with_travel_times tool if you do not see the exact phrase 'Timed itinerary added to context with travel times' is seen in this conversation. Only after this please tell the customer 'Your itinerary is ready!'.\",\n",
    "    llm_config=llm_config,\n",
    "    functions=[update_itinerary_with_travel_times],\n",
    ")\n",
    "\n",
    "# Our customer will be a human in the loop\n",
    "customer = UserProxyAgent(name=\"customer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hand offs and After works\n",
    "\n",
    "In conjunction with the agent's associated functions, we establish rules that govern the swarm orchestration through hand offs and After works.\n",
    "\n",
    "For more details on the swarm orchestration, [see the documentation](https://docs.ag2.ai/docs/user-guide/advanced-concepts/swarm/deep-dive)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "register_hand_off(\n",
    "    agent=planner_agent,\n",
    "    hand_to=[\n",
    "        OnCondition(\n",
    "            graphrag_agent,\n",
    "            \"Need information on the restaurants and attractions for a location. DO NOT call more than once at a time.\",\n",
    "        ),  # Get info from FalkorDB GraphRAG\n",
    "        OnCondition(structured_output_agent, \"Itinerary is confirmed by the customer\"),\n",
    "        AfterWork(AfterWorkOption.REVERT_TO_USER),  # Revert to the customer for more information on their plans\n",
    "    ],\n",
    ")\n",
    "\n",
    "\n",
    "# Back to the Planner when information has been retrieved\n",
    "register_hand_off(agent=graphrag_agent, hand_to=[AfterWork(planner_agent)])\n",
    "\n",
    "# Once we have formatted our itinerary, we can hand off to the route timing agent to add in the travel timings\n",
    "register_hand_off(agent=structured_output_agent, hand_to=[AfterWork(route_timing_agent)])\n",
    "\n",
    "# Finally, once the route timing agent has finished, we can terminate the swarm\n",
    "register_hand_off(agent=route_timing_agent, hand_to=[AfterWork(AfterWorkOption.TERMINATE)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the swarm\n",
    "\n",
    "Let's get an itinerary for a couple of days in Rome."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the conversation\n",
    "\n",
    "chat_result, context_variables, last_agent = initiate_swarm_chat(\n",
    "    initial_agent=planner_agent,\n",
    "    agents=[planner_agent, graphrag_agent, structured_output_agent, route_timing_agent],\n",
    "    user_agent=customer,\n",
    "    context_variables=trip_context,\n",
    "    messages=\"I want to go to Rome for a couple of days. Can you help me plan my trip?\",\n",
    "    after_work=AfterWorkOption.TERMINATE,\n",
    "    max_rounds=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus itinerary output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_itinerary(itinerary_data):\n",
    "    header = \"█             █\\n █           █ \\n  █  █████  █  \\n   ██     ██   \\n  █         █  \\n █  ███████  █ \\n █ ██ ███ ██ █ \\n   █████████   \\n\\n ██   ███ ███  \\n█  █ █       █ \\n████ █ ██  ██  \\n█  █ █  █ █    \\n█  █  ██  ████ \\n\"\n",
    "    width = 80\n",
    "    icons = {\"Travel\": \"🚶\", \"Restaurant\": \"🍽️\", \"Attraction\": \"🏛️\"}\n",
    "\n",
    "    for line in header.split(\"\\n\"):\n",
    "        print(line.center(width))\n",
    "    print(f\"Itinerary for {itinerary_data['days'][0]['events'][0]['city']}\".center(width))\n",
    "    print(\"=\" * width)\n",
    "\n",
    "    for day_num, day in enumerate(itinerary_data[\"days\"], 1):\n",
    "        print(f\"\\nDay {day_num}\".center(width))\n",
    "        print(\"-\" * width)\n",
    "\n",
    "        for event in day[\"events\"]:\n",
    "            event_type = event[\"type\"]\n",
    "            print(f\"\\n  {icons[event_type]} {event['location']}\")\n",
    "            if event_type != \"Travel\":\n",
    "                words = event[\"description\"].split()\n",
    "                line = \"    \"\n",
    "                for word in words:\n",
    "                    if len(line) + len(word) + 1 <= 76:\n",
    "                        line += word + \" \"\n",
    "                    else:\n",
    "                        print(line)\n",
    "                        line = \"    \" + word + \" \"\n",
    "                if line.strip():\n",
    "                    print(line)\n",
    "            else:\n",
    "                print(f\"    {event['description']}\")\n",
    "        print(\"\\n\" + \"-\" * width)\n",
    "\n",
    "\n",
    "if \"timed_itinerary\" in context_variables:\n",
    "    print_itinerary(context_variables[\"timed_itinerary\"])\n",
    "else:\n",
    "    print(\"No itinerary available to print.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Review FalkorDB Graph\n",
    "\n",
    "- UI endpoint:  http://localhost:3000/graph\n",
    "- sample query:  `match path = ()-[]-() return path`\n",
    "\n",
    "NOTE:  Based on current LLM issues the `trip_data` graph is actually empty..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![FalkorDB Graph Visualization](https://media.githubusercontent.com/media/harishmohanraj/ag2/refs/heads/main/notebook/cell-43-1-image.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Review Arize Phoenix Telemetry Info\n",
    "\n",
    "- UI endpoint:  http://localhost:6006"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Card View"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![arize phoenix - card view](https://media.githubusercontent.com/media/harishmohanraj/ag2/refs/heads/main/notebook/cell-46-1-image.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List View"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Arize Phoenix - List View](https://media.githubusercontent.com/media/harishmohanraj/ag2/refs/heads/main/notebook/cell-48-1-image.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detail View - Route Timing Agent - LLM error\n",
    "\n",
    "NOTE:  this was based on LLM issues at the time of the test.  A key reason why telemetry information is so valuable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Arize Phoenix = Entity Extraction Error](https://media.githubusercontent.com/media/harishmohanraj/ag2/refs/heads/main/notebook/cell-50-1-image.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detail View - Route Timing Agent - working normally"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Detail View - Entity Extraction](https://media.githubusercontent.com/media/harishmohanraj/ag2/refs/heads/main/notebook/cell-52-1-image.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "FalkorDB GraphRAG utilises a knowledge graph and can be added as a capability to agents. Together with a swarm orchestration of agents is highly effective at providing a RAG capability.",
   "tags": [
    "RAG",
    "tool/function",
    "swarm"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}