diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ae197c54e4..57db0b7fb4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -47,7 +47,8 @@ repos: website/docs/tutorial/code-executors.ipynb | website/docs/topics/code-execution/custom-executor.ipynb | website/docs/topics/non-openai-models/cloud-gemini.ipynb | - notebook/.* + notebook/.* | + test/agentchat/contrib/graph_rag/trip_planner_data/.* )$ # See https://jaredkhan.com/blog/mypy-pre-commit - repo: local diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py index 114e6d8ed3..ef984279d9 100644 --- a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py @@ -3,9 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 import os -from dataclasses import dataclass, field +import warnings from typing import List +from falkordb import FalkorDB, Graph from graphrag_sdk import KnowledgeGraph, Source from graphrag_sdk.model_config import KnowledgeGraphModelConfig from graphrag_sdk.models import GenerativeModel @@ -35,6 +36,8 @@ def __init__( Initialize a FalkorDB knowledge graph. Please also refer to https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/kg.py + TODO: Fix LLM API cost calculation for FalkorDB useages. + Args: name (str): Knowledge graph name. host (str): FalkorDB hostname. @@ -53,8 +56,38 @@ def __init__( self.model = model self.model_config = KnowledgeGraphModelConfig.with_model(model) self.ontology = ontology + self.knowledge_graph = None + self.falkordb = FalkorDB(host=self.host, port=self.port, username=self.username, password=self.password) + + def connect_db(self): + """ + Connect to an existing knowledge graph. + """ + if self.name in self.falkordb.list_graphs(): + try: + self.ontology = self._load_ontology_from_db(self.name) + except Exception: + warnings.warn("Graph Ontology is not loaded.") + + if self.ontology is None: + raise ValueError(f"Ontology of the knowledge graph '{self.name}' can't be None.") + + self.knowledge_graph = KnowledgeGraph( + name=self.name, + host=self.host, + port=self.port, + username=self.username, + password=self.password, + model_config=self.model_config, + ontology=self.ontology, + ) + + # Establishing a chat session will maintain the history + self._chat_session = self.knowledge_graph.chat_session() + else: + raise ValueError(f"Knowledge graph '{self.name}' does not exist") - def init_db(self, input_doc: List[Document] | None): + def init_db(self, input_doc: List[Document]): """ Build the knowledge graph with input documents. """ @@ -81,10 +114,16 @@ def init_db(self, input_doc: List[Document] | None): ontology=self.ontology, ) - # Establish a chat session, this will maintain the history - self._chat_session = self.knowledge_graph.chat_session() self.knowledge_graph.process_sources(sources) + # Establishing a chat session will maintain the history + self._chat_session = self.knowledge_graph.chat_session() + + # Save Ontology to graph for future access. + self._save_ontology_to_db(self.name, self.ontology) + else: + raise ValueError("No input documents could be loaded.") + def add_records(self, new_records: List) -> bool: raise NotImplementedError("This method is not supported by FalkorDB SDK yet.") @@ -101,7 +140,7 @@ def query(self, question: str, n_results: int = 1, **kwargs) -> GraphStoreQueryR Returns: FalkorGraphQueryResult """ if self.knowledge_graph is None: - raise ValueError("Knowledge graph is not created.") + raise ValueError("Knowledge graph has not been selected or created.") response = self._chat_session.send_message(question) @@ -109,3 +148,18 @@ def query(self, question: str, n_results: int = 1, **kwargs) -> GraphStoreQueryR self._chat_session.last_answer = response["response"] return GraphStoreQueryResult(answer=response["response"], results=[]) + + def __get_ontology_storage_graph(self, graph_name: str) -> Graph: + ontology_table_name = graph_name + "_ontology" + return self.falkordb.select_graph(ontology_table_name) + + def _save_ontology_to_db(self, graph_name: str, ontology: Ontology): + """ + Save graph ontology to a separate table with {graph_name}_ontology + """ + graph = self.__get_ontology_storage_graph(graph_name) + ontology.save_to_graph(graph) + + def _load_ontology_from_db(self, graph_name: str) -> Ontology: + graph = self.__get_ontology_storage_graph(graph_name) + return Ontology.from_graph(graph) diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py index 7e8bf9cc90..b5432403c5 100644 --- a/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py +++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py @@ -55,6 +55,8 @@ def _reply_using_falkordb_query( Query FalkorDB and return the message. Internally, it utilises OpenAI to generate a reply based on the given messages. The history with FalkorDB is also logged and updated. + The agent's system message will be incorporated into the query, if it's not blank. + If no results are found, a default message is returned: "I'm sorry, I don't have an answer for that." Args: @@ -66,16 +68,38 @@ def _reply_using_falkordb_query( Returns: A tuple containing a boolean indicating success and the assistant's reply. """ - question = self._get_last_question(messages[-1]) + # question = self._get_last_question(messages[-1]) + question = self._messages_summary(messages, recipient.system_message) result: GraphStoreQueryResult = self.query_engine.query(question) return True, result.answer if result.answer else "I'm sorry, I don't have an answer for that." - def _get_last_question(self, message: Union[Dict, str]): - """Retrieves the last message from the conversation history.""" - if isinstance(message, str): - return message - if isinstance(message, Dict): - if "content" in message: - return message["content"] - return None + def _messages_summary(self, messages: Union[Dict, str], system_message: str) -> str: + """Summarize the messages in the conversation history. Excluding any message with 'tool_calls' and 'tool_responses' + Includes the 'name' (if it exists) and the 'content', with a new line between each one, like: + customer: + + + agent: + + """ + + if isinstance(messages, str): + if system_message: + summary = f"IMPORTANT: {system_message}\nContext:\n\n{messages}" + else: + return messages + + elif isinstance(messages, List): + summary = "" + for message in messages: + if "content" in message and "tool_calls" not in message and "tool_responses" not in message: + summary += f"{message.get('name', '')}: {message.get('content','')}\n\n" + + if system_message: + summary = f"IMPORTANT: {system_message}\nContext:\n\n{summary}" + + return summary + + else: + raise ValueError("Invalid messages format. Must be a list of messages or a string.") diff --git a/notebook/agentchat_swarm_graphrag_trip_planner.ipynb b/notebook/agentchat_swarm_graphrag_trip_planner.ipynb new file mode 100644 index 0000000000..13a2763028 --- /dev/null +++ b/notebook/agentchat_swarm_graphrag_trip_planner.ipynb @@ -0,0 +1,1183 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Trip planning with a FalkorDB GraphRAG agent using a Swarm\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", + "- 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 activites\n", + "- Swarm orchestration utilising context variables" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "````{=mdx}\n", + ":::info Requirements\n", + "FalkorDB's GraphRAG-SDK is a dependency for this notebook, which can be installed with ag2 via pip:\n", + "\n", + "```bash\n", + "pip install ag2[graph_rag_falkor_db]\n", + "```\n", + "\n", + "For more information, please refer to the [installation guide](/docs/installation/).\n", + ":::\n", + "````\n", + "\n", + "### Pydantic\n", + "\n", + "Please ensure you have Pydantic version 2+ installed.\n", + "```bash\n", + "pip install --upgrade pydantic\n", + "```\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Running a FalkorDB\n", + "\n", + "**Note:** You need to have a FalkorDB graph database running. If you are running one in a Docker container, please ensure your Docker network is setup to allow access to it.\n", + "\n", + "In this example, we've set the FalkorDB host and port, please adjust them accordingly. For how to set up FalkorDB, please refer to [https://docs.falkordb.com/](https://docs.falkordb.com/).\n" + ] + }, + { + "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 your environment variable `GOOGLE_MAP_API_KEY` to the key" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Set Configuration and OpenAI API Key\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](/docs/topics/llm_configuration).\n", + ":::\n", + "````" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/usr/local/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + } + ], + "source": [ + "import os\n", + "\n", + "import autogen\n", + "\n", + "config_list = autogen.config_list_from_json(env_or_file=\"OAI_CONFIG_LIST\", filter_dict={\"model\": [\"gpt-4o\"]})\n", + "llm_config = {\"config_list\": config_list, \"timeout\": 120}\n", + "\n", + "# Put the OpenAI API key into the environment\n", + "os.environ[\"OPENAI_API_KEY\"] = config_list[0][\"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": 2, + "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.json\",\n", + " \"../test/agentchat/contrib/graph_rag/trip_planner_data/cities.json\",\n", + " \"../test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.json\",\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": 3, + "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": [ + "### Establish FalkorDB and load\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": 4, + "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=\"192.168.0.115\", # Change\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": 5, + "metadata": {}, + "outputs": [], + "source": [ + "# IMPORTS\n", + "import copy\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", + " AFTER_WORK,\n", + " ON_CONDITION,\n", + " AfterWorkOption,\n", + " SwarmAgent,\n", + " SwarmResult,\n", + " UserProxyAgent,\n", + " initiate_swarm_chat,\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": 6, + "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": 7, + "metadata": {}, + "outputs": [], + "source": [ + "def _fetch_travel_time(origin: str, destination: str) -> dict:\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", + " 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", + " \"\"\"\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": 8, + "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": 9, + "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", + "\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 SwarmAgents and a UserProxyAgent (human) which the swarm will interact with." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "# Planner agent, interacting with the customer and GraphRag agent, to create an itinerary\n", + "planner_agent = SwarmAgent(\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 = SwarmAgent(\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_config_list = copy.deepcopy(config_list)\n", + "for config in structured_config_list:\n", + " config[\"response_format\"] = Itinerary\n", + "\n", + "structured_output_agent = SwarmAgent(\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={\"config_list\": structured_config_list, \"timeout\": 120},\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 = SwarmAgent(\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://ag2ai.github.io/ag2/docs/topics/swarm)." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "planner_agent.register_hand_off(\n", + " hand_to=[\n", + " ON_CONDITION(\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", + " ON_CONDITION(structured_output_agent, \"Itinerary is confirmed by the customer\"),\n", + " AFTER_WORK(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", + "graphrag_agent.register_hand_off(hand_to=[AFTER_WORK(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", + "structured_output_agent.register_hand_off(hand_to=[AFTER_WORK(route_timing_agent)])\n", + "\n", + "# Finally, once the route timing agent has finished, we can terminate the swarm\n", + "route_timing_agent.register_hand_off(\n", + " hand_to=[AFTER_WORK(AfterWorkOption.TERMINATE)] # Once this agent has finished, the swarm can terminate\n", + ")" + ] + }, + { + "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": 12, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[33mcustomer\u001b[0m (to chat_manager):\n", + "\n", + "I want to go to Rome for a couple of days. Can you help me plan my trip?\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: planner_agent\n", + "\u001b[0m\n", + "\u001b[33mplanner_agent\u001b[0m (to chat_manager):\n", + "\n", + "Sure! How many days will you be staying in Rome? Additionally, what type of activities are you interested in doing, and do you have any specific preferences for meals (e.g., certain type of cuisine or dining experience)?\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: customer\n", + "\u001b[0m\n", + "\u001b[33mcustomer\u001b[0m (to chat_manager):\n", + "\n", + "2 days please, I want lunch and dinner, with an attraction to see in the morning and after lunch.\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: planner_agent\n", + "\u001b[0m\n", + "\u001b[33mplanner_agent\u001b[0m (to chat_manager):\n", + "\n", + "\u001b[32m***** Suggested tool call (call_vQMpso8aOomdfq8S2uCRlnzj): transfer_to_graphrag_agent *****\u001b[0m\n", + "Arguments: \n", + "{}\n", + "\u001b[32m*******************************************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: Tool_Execution\n", + "\u001b[0m\n", + "\u001b[35m\n", + ">>>>>>>> EXECUTING FUNCTION transfer_to_graphrag_agent...\u001b[0m\n", + "\u001b[33mTool_Execution\u001b[0m (to chat_manager):\n", + "\n", + "\u001b[32m***** Response from calling tool (call_vQMpso8aOomdfq8S2uCRlnzj) *****\u001b[0m\n", + "SwarmAgent --> graphrag_agent\n", + "\u001b[32m**********************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: graphrag_agent\n", + "\u001b[0m\n", + "\u001b[33mgraphrag_agent\u001b[0m (to chat_manager):\n", + "\n", + "Restaurants in Rome:\n", + "1. Trattoria da Enzo - A cozy trattoria known for its traditional Roman dishes and welcoming atmosphere (Rating: 4.5).\n", + "2. Il Pagliaccio - An elegant Michelin-starred restaurant offering contemporary Italian cuisine (Rating: 4.8).\n", + "3. Tonnarello - Casual stop for pasta, meatballs & other simple Roman dishes, plus patio seating & acoustic guitar (Rating: 4.7).\n", + "4. Osteria delle Commari - Local home style Roman restaurant near Vatican Museums (Rating: 4.3).\n", + "5. Armando al Pantheon - Traditional Roman cuisine since 1961 near the Pantheon (Rating: 4.8).\n", + "6. Salumeria Roscioli Ristorante - Known for pasta alla gricia, amatriciana, and an impressive wine cellar (Rating: 4.5).\n", + "7. L'Arcangelo - Known for some of the best supplì in the city and Thursday gnocchi (Rating: 4.8).\n", + "\n", + "Attractions in Rome:\n", + "1. Colosseum - An ancient amphitheater known for gladiatorial contests and public spectacles.\n", + "2. Vatican Museums - A complex of museums and galleries showcasing works of art collected by Popes over centuries.\n", + "3. Trevi Fountain - A Baroque fountain known for its stunning sculptures and tradition of tossing coins.\n", + "4. Basilica di Santa Maria Maggiore - A major basilica known for its rich history and impressive architecture.\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: planner_agent\n", + "\u001b[0m\n", + "\u001b[33mplanner_agent\u001b[0m (to chat_manager):\n", + "\n", + "Great! Here's a suggested itinerary for your 2-day trip to Rome:\n", + "\n", + "### Day 1:\n", + "\n", + "- **Morning Attraction:**\n", + " - **Type:** Attraction\n", + " - **Location:** Colosseum\n", + " - **City:** Rome\n", + " - **Description:** Explore the ancient amphitheater known for gladiatorial contests and public spectacles.\n", + "\n", + "- **Lunch:**\n", + " - **Type:** Restaurant\n", + " - **Location:** Trattoria da Enzo\n", + " - **City:** Rome\n", + " - **Description:** Enjoy traditional Roman dishes in a cozy and welcoming atmosphere.\n", + "\n", + "- **Afternoon Attraction:**\n", + " - **Type:** Attraction\n", + " - **Location:** Vatican Museums\n", + " - **City:** Rome\n", + " - **Description:** Visit a complex of museums and galleries showcasing works of art collected by Popes over centuries.\n", + "\n", + "- **Dinner:**\n", + " - **Type:** Restaurant\n", + " - **Location:** Armando al Pantheon\n", + " - **City:** Rome\n", + " - **Description:** Experience traditional Roman cuisine since 1961, located near the Pantheon.\n", + "\n", + "\n", + "### Day 2:\n", + "\n", + "- **Morning Attraction:**\n", + " - **Type:** Attraction\n", + " - **Location:** Trevi Fountain\n", + " - **City:** Rome\n", + " - **Description:** Marvel at the stunning sculptures of the Baroque fountain and join the tradition of tossing coins.\n", + "\n", + "- **Lunch:**\n", + " - **Type:** Restaurant\n", + " - **Location:** Tonnarello\n", + " - **City:** Rome\n", + " - **Description:** Relish in simple Roman dishes such as pasta and meatballs in a casual setting with patio seating.\n", + "\n", + "- **Afternoon Attraction:**\n", + " - **Type:** Attraction\n", + " - **Location:** Basilica di Santa Maria Maggiore\n", + " - **City:** Rome\n", + " - **Description:** Visit a major basilica known for its rich history and impressive architecture.\n", + "\n", + "- **Dinner:**\n", + " - **Type:** Restaurant\n", + " - **Location:** L'Arcangelo\n", + " - **City:** Rome\n", + " - **Description:** Savor some of the best supplì in the city and enjoy traditional gnocchi on Thursdays.\n", + "\n", + "Let me know if you are happy with this itinerary or if there are any changes you would like to make!\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: customer\n", + "\u001b[0m\n", + "\u001b[33mcustomer\u001b[0m (to chat_manager):\n", + "\n", + "I don't feel like eating at L'Arcangelo, please change that for another option.\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: planner_agent\n", + "\u001b[0m\n", + "\u001b[33mplanner_agent\u001b[0m (to chat_manager):\n", + "\n", + "Certainly! Here's the updated itinerary for Day 2:\n", + "\n", + "### Day 2:\n", + "\n", + "- **Morning Attraction:**\n", + " - **Type:** Attraction\n", + " - **Location:** Trevi Fountain\n", + " - **City:** Rome\n", + " - **Description:** Marvel at the stunning sculptures of the Baroque fountain and join the tradition of tossing coins.\n", + "\n", + "- **Lunch:**\n", + " - **Type:** Restaurant\n", + " - **Location:** Tonnarello\n", + " - **City:** Rome\n", + " - **Description:** Relish in simple Roman dishes such as pasta and meatballs in a casual setting with patio seating.\n", + "\n", + "- **Afternoon Attraction:**\n", + " - **Type:** Attraction\n", + " - **Location:** Basilica di Santa Maria Maggiore\n", + " - **City:** Rome\n", + " - **Description:** Visit a major basilica known for its rich history and impressive architecture.\n", + "\n", + "- **Dinner:**\n", + " - **Type:** Restaurant\n", + " - **Location:** Il Pagliaccio\n", + " - **City:** Rome\n", + " - **Description:** Indulge in elegant, contemporary Italian cuisine at a Michelin-starred restaurant.\n", + "\n", + "Let me know if this itinerary meets your preferences!\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: customer\n", + "\u001b[0m\n", + "\u001b[33mcustomer\u001b[0m (to chat_manager):\n", + "\n", + "Can you research and show me all the available restaurant options.\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: planner_agent\n", + "\u001b[0m\n", + "\u001b[33mplanner_agent\u001b[0m (to chat_manager):\n", + "\n", + "Here are the available restaurant options for your trip to Rome:\n", + "\n", + "1. **Trattoria da Enzo**\n", + " - **Description:** A cozy trattoria known for its traditional Roman dishes and welcoming atmosphere.\n", + " - **Rating:** 4.5\n", + "\n", + "2. **Il Pagliaccio**\n", + " - **Description:** An elegant Michelin-starred restaurant offering contemporary Italian cuisine.\n", + " - **Rating:** 4.8\n", + "\n", + "3. **Tonnarello**\n", + " - **Description:** Casual stop for pasta, meatballs & other simple Roman dishes, plus patio seating & acoustic guitar.\n", + " - **Rating:** 4.7\n", + "\n", + "4. **Osteria delle Commari**\n", + " - **Description:** Local home style Roman restaurant near Vatican Museums.\n", + " - **Rating:** 4.3\n", + "\n", + "5. **Armando al Pantheon**\n", + " - **Description:** Traditional Roman cuisine since 1961 near the Pantheon.\n", + " - **Rating:** 4.8\n", + "\n", + "6. **Salumeria Roscioli Ristorante**\n", + " - **Description:** Known for pasta alla gricia, amatriciana, and an impressive wine cellar.\n", + " - **Rating:** 4.5\n", + "\n", + "7. **L'Arcangelo**\n", + " - **Description:** Known for some of the best supplì in the city and Thursday gnocchi.\n", + " - **Rating:** 4.8\n", + "\n", + "Please let me know which restaurants you'd like to include in your itinerary or if there are any changes you'd like to make!\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: customer\n", + "\u001b[0m\n", + "\u001b[33mcustomer\u001b[0m (to chat_manager):\n", + "\n", + "Tonnarello's please.\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: planner_agent\n", + "\u001b[0m\n", + "\u001b[33mplanner_agent\u001b[0m (to chat_manager):\n", + "\n", + "\u001b[32m***** Suggested tool call (call_sROJ1YfU63SDCRfUbRuGqjVa): mark_itinerary_as_complete *****\u001b[0m\n", + "Arguments: \n", + "{\"final_itinerary\": \"### Day 1:\\n\\n- **Morning Attraction:**\\n - **Type:** Attraction\\n - **Location:** Colosseum\\n - **City:** Rome\\n - **Description:** Explore the ancient amphitheater known for gladiatorial contests and public spectacles.\\n\\n- **Lunch:**\\n - **Type:** Restaurant\\n - **Location:** Trattoria da Enzo\\n - **City:** Rome\\n - **Description:** Enjoy traditional Roman dishes in a cozy and welcoming atmosphere.\\n\\n- **Afternoon Attraction:**\\n - **Type:** Attraction\\n - **Location:** Vatican Museums\\n - **City:** Rome\\n - **Description:** Visit a complex of museums and galleries showcasing works of art collected by Popes over centuries.\\n\\n- **Dinner:**\\n - **Type:** Restaurant\\n - **Location:** Armando al Pantheon\\n - **City:** Rome\\n - **Description:** Experience traditional Roman cuisine since 1961, located near the Pantheon.\\n\\n\\n### Day 2:\\n\\n- **Morning Attraction:**\\n - **Type:** Attraction\\n - **Location:** Trevi Fountain\\n - **City:** Rome\\n - **Description:** Marvel at the stunning sculptures of the Baroque fountain and join the tradition of tossing coins.\\n\\n- **Lunch:**\\n - **Type:** Restaurant\\n - **Location:** Tonnarello\\n - **City:** Rome\\n - **Description:** Relish in simple Roman dishes such as pasta and meatballs in a casual setting with patio seating.\\n\\n- **Afternoon Attraction:**\\n - **Type:** Attraction\\n - **Location:** Basilica di Santa Maria Maggiore\\n - **City:** Rome\\n - **Description:** Visit a major basilica known for its rich history and impressive architecture.\\n\\n- **Dinner:**\\n - **Type:** Restaurant\\n - **Location:** Tonnarello\\n - **City:** Rome\\n - **Description:** Return to Tonnarello for a delightful dining experience with a variety of Roman specialties.\"}\n", + "\u001b[32m*******************************************************************************************\u001b[0m\n", + "\u001b[32m***** Suggested tool call (call_NBw71N4pS66h8VLlgu5nvveN): transfer_to_structured_output_agent *****\u001b[0m\n", + "Arguments: \n", + "{}\n", + "\u001b[32m****************************************************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: Tool_Execution\n", + "\u001b[0m\n", + "\u001b[35m\n", + ">>>>>>>> EXECUTING FUNCTION mark_itinerary_as_complete...\u001b[0m\n", + "\u001b[35m\n", + ">>>>>>>> EXECUTING FUNCTION transfer_to_structured_output_agent...\u001b[0m\n", + "\u001b[33mTool_Execution\u001b[0m (to chat_manager):\n", + "\n", + "\u001b[32m***** Response from calling tool (call_sROJ1YfU63SDCRfUbRuGqjVa) *****\u001b[0m\n", + "Itinerary recorded and confirmed.\n", + "\u001b[32m**********************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m***** Response from calling tool (call_NBw71N4pS66h8VLlgu5nvveN) *****\u001b[0m\n", + "SwarmAgent --> structured_output_agent\n", + "\u001b[32m**********************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: structured_output_agent\n", + "\u001b[0m\n", + "\u001b[33mstructured_output_agent\u001b[0m (to chat_manager):\n", + "\n", + "{\"days\":[{\"events\":[{\"type\":\"Attraction\",\"location\":\"Colosseum\",\"city\":\"Rome\",\"description\":\"Explore the ancient amphitheater known for gladiatorial contests and public spectacles.\"},{\"type\":\"Restaurant\",\"location\":\"Trattoria da Enzo\",\"city\":\"Rome\",\"description\":\"Enjoy traditional Roman dishes in a cozy and welcoming atmosphere.\"},{\"type\":\"Attraction\",\"location\":\"Vatican Museums\",\"city\":\"Rome\",\"description\":\"Visit a complex of museums and galleries showcasing works of art collected by Popes over centuries.\"},{\"type\":\"Restaurant\",\"location\":\"Armando al Pantheon\",\"city\":\"Rome\",\"description\":\"Experience traditional Roman cuisine since 1961, located near the Pantheon.\"}]},{\"events\":[{\"type\":\"Attraction\",\"location\":\"Trevi Fountain\",\"city\":\"Rome\",\"description\":\"Marvel at the stunning sculptures of the Baroque fountain and join the tradition of tossing coins.\"},{\"type\":\"Restaurant\",\"location\":\"Tonnarello\",\"city\":\"Rome\",\"description\":\"Relish in simple Roman dishes such as pasta and meatballs in a casual setting with patio seating.\"},{\"type\":\"Attraction\",\"location\":\"Basilica di Santa Maria Maggiore\",\"city\":\"Rome\",\"description\":\"Visit a major basilica known for its rich history and impressive architecture.\"},{\"type\":\"Restaurant\",\"location\":\"Tonnarello\",\"city\":\"Rome\",\"description\":\"Return to Tonnarello for a delightful dining experience with a variety of Roman specialties.\"}]}]}\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: route_timing_agent\n", + "\u001b[0m\n", + "\u001b[33mroute_timing_agent\u001b[0m (to chat_manager):\n", + "\n", + "None\n", + "\u001b[32m***** Suggested tool call (call_6nnUKnD3VveszMT8OvkqUTCC): update_itinerary_with_travel_times *****\u001b[0m\n", + "Arguments: \n", + "{}\n", + "\u001b[32m***************************************************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: Tool_Execution\n", + "\u001b[0m\n", + "\u001b[35m\n", + ">>>>>>>> EXECUTING FUNCTION update_itinerary_with_travel_times...\u001b[0m\n", + "\u001b[33mTool_Execution\u001b[0m (to chat_manager):\n", + "\n", + "\u001b[32m***** Response from calling tool (call_6nnUKnD3VveszMT8OvkqUTCC) *****\u001b[0m\n", + "Structured itinerary not found, please create the structured output, structured_output_agent.\n", + "\u001b[32m**********************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: structured_output_agent\n", + "\u001b[0m\n", + "\u001b[33mstructured_output_agent\u001b[0m (to chat_manager):\n", + "\n", + "\u001b[32m***** Suggested tool call (call_uTWKi1KdcsaFK9RxAeFUIWbU): create_structured_itinerary *****\u001b[0m\n", + "Arguments: \n", + "{\"structured_itinerary\":\"{\\\"days\\\":[{\\\"events\\\":[{\\\"type\\\":\\\"Attraction\\\",\\\"location\\\":\\\"Colosseum\\\",\\\"city\\\":\\\"Rome\\\",\\\"description\\\":\\\"Explore the ancient amphitheater known for gladiatorial contests and public spectacles.\\\"},{\\\"type\\\":\\\"Restaurant\\\",\\\"location\\\":\\\"Trattoria da Enzo\\\",\\\"city\\\":\\\"Rome\\\",\\\"description\\\":\\\"Enjoy traditional Roman dishes in a cozy and welcoming atmosphere.\\\"},{\\\"type\\\":\\\"Attraction\\\",\\\"location\\\":\\\"Vatican Museums\\\",\\\"city\\\":\\\"Rome\\\",\\\"description\\\":\\\"Visit a complex of museums and galleries showcasing works of art collected by Popes over centuries.\\\"},{\\\"type\\\":\\\"Restaurant\\\",\\\"location\\\":\\\"Armando al Pantheon\\\",\\\"city\\\":\\\"Rome\\\",\\\"description\\\":\\\"Experience traditional Roman cuisine since 1961, located near the Pantheon.\\\"}]},{\\\"events\\\":[{\\\"type\\\":\\\"Attraction\\\",\\\"location\\\":\\\"Trevi Fountain\\\",\\\"city\\\":\\\"Rome\\\",\\\"description\\\":\\\"Marvel at the stunning sculptures of the Baroque fountain and join the tradition of tossing coins.\\\"},{\\\"type\\\":\\\"Restaurant\\\",\\\"location\\\":\\\"Tonnarello\\\",\\\"city\\\":\\\"Rome\\\",\\\"description\\\":\\\"Relish in simple Roman dishes such as pasta and meatballs in a casual setting with patio seating.\\\"},{\\\"type\\\":\\\"Attraction\\\",\\\"location\\\":\\\"Basilica di Santa Maria Maggiore\\\",\\\"city\\\":\\\"Rome\\\",\\\"description\\\":\\\"Visit a major basilica known for its rich history and impressive architecture.\\\"},{\\\"type\\\":\\\"Restaurant\\\",\\\"location\\\":\\\"Tonnarello\\\",\\\"city\\\":\\\"Rome\\\",\\\"description\\\":\\\"Return to Tonnarello for a delightful dining experience with a variety of Roman specialties.\\\"}]}]}\"}\n", + "\u001b[32m********************************************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: Tool_Execution\n", + "\u001b[0m\n", + "\u001b[35m\n", + ">>>>>>>> EXECUTING FUNCTION create_structured_itinerary...\u001b[0m\n", + "\u001b[33mTool_Execution\u001b[0m (to chat_manager):\n", + "\n", + "\u001b[32m***** Response from calling tool (call_uTWKi1KdcsaFK9RxAeFUIWbU) *****\u001b[0m\n", + "Structured itinerary stored.\n", + "\u001b[32m**********************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: route_timing_agent\n", + "\u001b[0m\n", + "\u001b[33mroute_timing_agent\u001b[0m (to chat_manager):\n", + "\n", + "\u001b[32m***** Suggested tool call (call_7alJoHYs74wsPIz4WhMrj7Zu): update_itinerary_with_travel_times *****\u001b[0m\n", + "Arguments: \n", + "{}\n", + "\u001b[32m***************************************************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: Tool_Execution\n", + "\u001b[0m\n", + "\u001b[35m\n", + ">>>>>>>> EXECUTING FUNCTION update_itinerary_with_travel_times...\u001b[0m\n", + "\u001b[33mTool_Execution\u001b[0m (to chat_manager):\n", + "\n", + "\u001b[32m***** Response from calling tool (call_7alJoHYs74wsPIz4WhMrj7Zu) *****\u001b[0m\n", + "Timed itinerary added to context with travel times\n", + "\u001b[32m**********************************************************************\u001b[0m\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32m\n", + "Next speaker: route_timing_agent\n", + "\u001b[0m\n", + "\u001b[33mroute_timing_agent\u001b[0m (to chat_manager):\n", + "\n", + "Your itinerary is ready!\n", + "\n", + "--------------------------------------------------------------------------------\n" + ] + } + ], + "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": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " █ █ \n", + " █ █ \n", + " █ █████ █ \n", + " ██ ██ \n", + " █ █ \n", + " █ ███████ █ \n", + " █ ██ ███ ██ █ \n", + " █████████ \n", + " \n", + " ██ ███ ███ \n", + " █ █ █ █ \n", + " ████ █ ██ ██ \n", + " █ █ █ █ █ \n", + " █ █ ██ ████ \n", + " \n", + " Itinerary for Rome \n", + "================================================================================\n", + " \n", + "Day 1 \n", + "--------------------------------------------------------------------------------\n", + "\n", + " 🏛️ Colosseum\n", + " Explore the ancient amphitheater known for gladiatorial contests and \n", + " public spectacles. \n", + "\n", + " 🚶 walking from Colosseum to Trattoria da Enzo\n", + " 24 mins, (1.8 km)\n", + "\n", + " 🍽️ Trattoria da Enzo\n", + " Enjoy traditional Roman dishes in a cozy and welcoming atmosphere. \n", + "\n", + " 🚶 walking from Trattoria da Enzo to Vatican Museums\n", + " 50 mins, (3.5 km)\n", + "\n", + " 🏛️ Vatican Museums\n", + " Visit a complex of museums and galleries showcasing works of art \n", + " collected by Popes over centuries. \n", + "\n", + " 🚶 walking from Vatican Museums to Armando al Pantheon\n", + " 36 mins, (2.6 km)\n", + "\n", + " 🍽️ Armando al Pantheon\n", + " Experience traditional Roman cuisine since 1961, located near the \n", + " Pantheon. \n", + "\n", + "--------------------------------------------------------------------------------\n", + " \n", + "Day 2 \n", + "--------------------------------------------------------------------------------\n", + "\n", + " 🏛️ Trevi Fountain\n", + " Marvel at the stunning sculptures of the Baroque fountain and join the \n", + " tradition of tossing coins. \n", + "\n", + " 🚶 walking from Trevi Fountain to Tonnarello\n", + " 31 mins, (2.2 km)\n", + "\n", + " 🍽️ Tonnarello\n", + " Relish in simple Roman dishes such as pasta and meatballs in a casual \n", + " setting with patio seating. \n", + "\n", + " 🚶 walking from Tonnarello to Basilica di Santa Maria Maggiore\n", + " 46 mins, (3.2 km)\n", + "\n", + " 🏛️ Basilica di Santa Maria Maggiore\n", + " Visit a major basilica known for its rich history and impressive \n", + " architecture. \n", + "\n", + " 🚶 walking from Basilica di Santa Maria Maggiore to Tonnarello\n", + " 44 mins, (3.2 km)\n", + "\n", + " 🍽️ Tonnarello\n", + " Return to Tonnarello for a delightful dining experience with a variety \n", + " of Roman specialties. \n", + "\n", + "--------------------------------------------------------------------------------\n" + ] + } + ], + "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.\")" + ] + } + ], + "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.10" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/setup.py b/setup.py index d8b1446932..df1d4f51d4 100644 --- a/setup.py +++ b/setup.py @@ -60,9 +60,7 @@ retrieve_chat_pgvector = [*retrieve_chat, "pgvector>=0.2.5"] -graph_rag_falkor_db = [ - "graphrag_sdk==0.3.3", -] +graph_rag_falkor_db = ["graphrag_sdk==0.3.3", "falkordb>=1.0.10"] if current_os in ["Windows", "Darwin"]: retrieve_chat_pgvector.extend(["psycopg[binary]>=3.1.18"]) diff --git a/test/agentchat/contrib/graph_rag/trip_planner_data/attractions.json b/test/agentchat/contrib/graph_rag/trip_planner_data/attractions.json new file mode 100644 index 0000000000..c929eb4c0e --- /dev/null +++ b/test/agentchat/contrib/graph_rag/trip_planner_data/attractions.json @@ -0,0 +1,142 @@ +[ + { + "id": 1, + "name": "Colosseum", + "description": "An ancient amphitheater known for gladiatorial contests and public spectacles.", + "city": "Rome", + "type": "Historical" + }, + { + "id": 2, + "name": "Vatican Museums", + "description": "A complex of museums and galleries showcasing works of art collected by Popes over centuries.", + "city": "Rome", + "type": "Art" + }, + { + "id": 3, + "name": "Trevi Fountain", + "description": "A Baroque fountain known for its stunning sculptures and tradition of tossing coins.", + "city": "Rome", + "type": "Landmark" + }, + { + "id": 4, + "name": "Duomo", + "description": "Florence's main cathedral, known for its red-tiled dome designed by Brunelleschi.", + "city": "Florence", + "type": "Religious" + }, + { + "id": 5, + "name": "Uffizi Gallery", + "description": "A prominent art museum featuring works by Botticelli, Michelangelo, and da Vinci.", + "city": "Florence", + "type": "Art" + }, + { + "id": 6, + "name": "Ponte Vecchio", + "description": "A medieval stone bridge over the Arno River, lined with shops selling jewelry and art.", + "city": "Florence", + "type": "Landmark" + }, + { + "id": 7, + "name": "Grand Canal", + "description": "The main waterway of Venice, lined with stunning palaces and bridges.", + "city": "Venice", + "type": "Scenic" + }, + { + "id": 8, + "name": "St. Mark's Basilica", + "description": "Venice's famous cathedral known for its Byzantine architecture and gilded mosaics.", + "city": "Venice", + "type": "Religious" + }, + { + "id": 9, + "name": "Rialto Bridge", + "description": "One of Venice's iconic bridges crossing the Grand Canal, offering beautiful views.", + "city": "Venice", + "type": "Landmark" + }, + { + "id": 10, + "name": "Mole Antonelliana", + "description": "A landmark building in Turin known for its towering spire and as the Museum of Cinema.", + "city": "Turin", + "type": "Landmark" + }, + { + "id": 11, + "name": "Royal Palace of Naples", + "description": "A historic palace in Naples, once home to the Bourbon Kings of Naples and Sicily.", + "city": "Naples", + "type": "Historical" + }, + { + "id": 12, + "name": "Teatro San Carlo", + "description": "The oldest continuously active opera house in the world, located in Naples.", + "city": "Naples", + "type": "Cultural" + }, + { + "id": 13, + "name": "Teatro alla Scala", + "description": "One of the world's most famous opera houses, located in Milan.", + "city": "Milan", + "type": "Cultural" + }, + { + "id": 14, + "name": "Galleria Vittorio Emanuele II", + "description": "A historic shopping arcade and a major landmark in Milan, known for its stunning architecture.", + "city": "Milan", + "type": "Landmark" + }, + { + "id": 15, + "name": "Teatro Regio", + "description": "A prominent opera house in Turin, known for its rich history and performances.", + "city": "Turin", + "type": "Cultural" + }, + { + "id": 16, + "name": "Palazzo dei Normanni", + "description": "A royal palace in Palermo, originally built in the 9th century.", + "city": "Palermo", + "type": "Historical" + }, + { + "id": 17, + "name": "Cattedrale di Palermo", + "description": "The cathedral church of the Roman Catholic Archdiocese of Palermo, dedicated to the Assumption of the Virgin Mary.", + "city": "Palermo", + "type": "Religious" + }, + { + "id": 18, + "name": "San Giovanni degli Eremiti", + "description": "A church in Palermo, renowned for its distinctive red domes.", + "city": "Palermo", + "type": "Religious" + }, + { + "id": 19, + "name": "San Lorenzo Maggiore", + "description": "A church in Naples, one of the oldest in the city, featuring Roman and early Christian remains.", + "city": "Naples", + "type": "Religious" + }, + { + "id": 20, + "name": "Basilica di Santa Maria Maggiore", + "description": "A major basilica in Rome, known for its rich history and impressive architecture.", + "city": "Rome", + "type": "Religious" + } +] diff --git a/test/agentchat/contrib/graph_rag/trip_planner_data/cities.json b/test/agentchat/contrib/graph_rag/trip_planner_data/cities.json new file mode 100644 index 0000000000..46eb4a98a0 --- /dev/null +++ b/test/agentchat/contrib/graph_rag/trip_planner_data/cities.json @@ -0,0 +1,72 @@ +[ + { + "id": 1, + "name": "Rome", + "country": "Italy", + "population": 2800000, + "weather": "Mediterranean" + }, + { + "id": 2, + "name": "Milan", + "country": "Italy", + "population": 1370000, + "weather": "Humid subtropical" + }, + { + "id": 3, + "name": "Naples", + "country": "Italy", + "population": 967000, + "weather": "Mediterranean" + }, + { + "id": 4, + "name": "Turin", + "country": "Italy", + "population": 875000, + "weather": "Humid subtropical" + }, + { + "id": 5, + "name": "Palermo", + "country": "Italy", + "population": 673000, + "weather": "Mediterranean" + }, + { + "id": 6, + "name": "Genoa", + "country": "Italy", + "population": 580000, + "weather": "Mediterranean" + }, + { + "id": 7, + "name": "Bologna", + "country": "Italy", + "population": 390000, + "weather": "Humid subtropical" + }, + { + "id": 8, + "name": "Florence", + "country": "Italy", + "population": 383000, + "weather": "Temperate" + }, + { + "id": 9, + "name": "Venice", + "country": "Italy", + "population": 260000, + "weather": "Humid subtropical" + }, + { + "id": 10, + "name": "Verona", + "country": "Italy", + "population": 258000, + "weather": "Humid subtropical" + } +] diff --git a/test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.json b/test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.json new file mode 100644 index 0000000000..9500173b31 --- /dev/null +++ b/test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.json @@ -0,0 +1,98 @@ +[ + { + "id": 1, + "name": "Trattoria da Enzo", + "description": "A cozy trattoria known for its traditional Roman dishes and welcoming atmosphere.", + "city": "Rome", + "rating": 4.5, + "food_type": "Italian" + }, + { + "id": 2, + "name": "Il Pagliaccio", + "description": "An elegant Michelin-starred restaurant offering contemporary Italian cuisine.", + "city": "Rome", + "rating": 4.8, + "food_type": "Italian" + }, + { + "id": 3, + "name": "Bros' Burger", + "description": "A popular burger joint in Milan known for its gourmet burgers and relaxed vibe.", + "city": "Milan", + "rating": 4.3, + "food_type": "Burgers" + }, + { + "id": 4, + "name": "Osteria Francescana", + "description": "A world-renowned restaurant in Modena, known for its innovative Italian cuisine.", + "city": "Modena", + "rating": 4.9, + "food_type": "Italian" + }, + { + "id": 5, + "name": "La Bottega del Buon Caffè", + "description": "An upscale restaurant in Florence offering seasonal Tuscan cuisine with a modern twist.", + "city": "Florence", + "rating": 4.7, + "food_type": "Italian" + }, + { + "id": 6, + "name": "Antiche Carampane", + "description": "A charming trattoria in Venice known for its seafood dishes and Venetian specialties.", + "city": "Venice", + "rating": 4.6, + "food_type": "Seafood" + }, + { + "id": 7, + "name": "Ristorante Quadri", + "description": "A historic restaurant in Venice's Piazza San Marco, offering gourmet Italian cuisine.", + "city": "Venice", + "rating": 4.8, + "food_type": "Italian" + }, + { + "id": 8, + "name": "Da Cimino", + "description": "A family-run trattoria in Palermo known for its authentic Sicilian cuisine.", + "city": "Palermo", + "rating": 4.4, + "food_type": "Sicilian" + }, + { + "id": 9, + "name": "Il Gambero Rosso", + "description": "A renowned seafood restaurant in Naples, offering fresh catch of the day and traditional dishes.", + "city": "Naples", + "rating": 4.5, + "food_type": "Seafood" + }, + { + "id": 10, + "name": "Taverna del Capitano", + "description": "A charming restaurant in Sorrento known for its seafood and stunning views of the Bay of Naples.", + "city": "Sorrento", + "rating": 4.7, + "food_type": "Seafood" + }, + { + "id": 11, + "name": "Tonnarello", + "description": "Casual stop for pasta, meatballs & other simple Roman dishes, plus patio seating & acoustic guitar.", + "city": "Rome", + "rating": 4.7, + "food_type": "Italian" + }, + { + "id": 12, + "name": "Osteria delle Commari", + "description": "Local home style Roman restaurant near Vatican Museums.", + "city": "Rome", + "rating": 4.3, + "food_type": "Italian" + } +] diff --git a/website/docs/Examples.md b/website/docs/Examples.md index e1f4884e5c..c4e0bce998 100644 --- a/website/docs/Examples.md +++ b/website/docs/Examples.md @@ -36,12 +36,19 @@ Links to notebook examples: - OptiGuide for Solving a Supply Chain Optimization Problem with Nested Chats with a Coding Agent and a Safeguard Agent - [View Notebook](/docs/notebooks/agentchat_nestedchat_optiguide) - Conversational Chess with Nested Chats and Tool Use - [View Notebook](/docs/notebooks/agentchat_nested_chats_chess) +### Swarms +- Orchestrating agents in a Swarm - [View Notebook](/docs/notebooks/agentchat_swarm) + ### Applications - Automated Continual Learning from New Data - [View Notebook](/docs/notebooks/agentchat_stream) - [AutoAnny](https://github.com/ag2ai/build-with-ag2/tree/main/samples/apps/auto-anny) - A Discord bot built using AutoGen +### RAG + +- GraphRAG agent using FalkorDB (feat. swarms and Google Maps API) - [View Notebook](/docs/notebooks/agentchat_swarm_graphrag_trip_planner) + ### Tool Use - **Web Search**: Solve Tasks Requiring Web Info - [View Notebook](/docs/notebooks/agentchat_web_info) diff --git a/website/docs/topics/swarm.ipynb b/website/docs/topics/swarm.ipynb index 7e8263f0db..1724eea982 100644 --- a/website/docs/topics/swarm.ipynb +++ b/website/docs/topics/swarm.ipynb @@ -247,7 +247,7 @@ "\n", "\n", "# 1. A function that returns a value of \"success\" and updates the context variable \"1\" to True\n", - "def update_context_1(context_variables: dict) -> str:\n", + "def update_context_1(context_variables: dict) -> SwarmResult:\n", " context_variables[\"1\"] = True\n", " return SwarmResult(value=\"success\", context_variables=context_variables)\n", "\n", @@ -259,7 +259,7 @@ "\n", "\n", "# 3. A function that returns the value of \"success\", updates the context variable and transfers to agent 3\n", - "def update_context_2_and_transfer_to_3(context_variables: dict) -> str:\n", + "def update_context_2_and_transfer_to_3(context_variables: dict) -> SwarmResult:\n", " context_variables[\"2\"] = True\n", " return SwarmResult(value=\"success\", context_variables=context_variables, agent=agent_3)\n", "\n", @@ -269,7 +269,7 @@ " return random.randint(1, 100)\n", "\n", "\n", - "def update_context_3_with_random_number(context_variables: dict, random_number: int) -> str:\n", + "def update_context_3_with_random_number(context_variables: dict, random_number: int) -> SwarmResult:\n", " context_variables[\"3\"] = random_number\n", " return SwarmResult(value=\"success\", context_variables=context_variables)\n", "\n",