graphiti/examples/langgraph-agent/agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a ShoeBot Sales Agent using LangGraph and Graphiti\n",
    "\n",
    "The following example demonstrates building an agent using LangGraph. Graphiti is used to personalize agent responses based on information learned from prior conversations. Additionally, a database of products is loaded into the Graphiti graph, enabling the agent to speak to these products.\n",
    "\n",
    "The agent implements:\n",
    "- persistance of new chat turns to Graphiti and recall of relevant Facts using the most recent message.\n",
    "- a tool for querying Graphiti for shoe information\n",
    "- an in-memory MemorySaver to maintain agent state.\n",
    "\n",
    "## Install dependencies\n",
    "```shell\n",
    "pip install graphiti-core langchain-openai langgraph ipywidgets\n",
    "```\n",
    "\n",
    "Ensure that you've followed the Graphiti installation instructions. In particular, installation of `neo4j`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import asyncio\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import sys\n",
    "import uuid\n",
    "from contextlib import suppress\n",
    "from datetime import datetime\n",
    "from pathlib import Path\n",
    "from typing import Annotated\n",
    "\n",
    "import ipywidgets as widgets\n",
    "from dotenv import load_dotenv\n",
    "from IPython.display import Image, display\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_logging():\n",
    "    logger = logging.getLogger()\n",
    "    logger.setLevel(logging.ERROR)\n",
    "    console_handler = logging.StreamHandler(sys.stdout)\n",
    "    console_handler.setLevel(logging.INFO)\n",
    "    formatter = logging.Formatter('%(name)s - %(levelname)s - %(message)s')\n",
    "    console_handler.setFormatter(formatter)\n",
    "    logger.addHandler(console_handler)\n",
    "    return logger\n",
    "\n",
    "\n",
    "logger = setup_logging()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LangSmith integration (Optional)\n",
    "\n",
    "If you'd like to trace your agent using LangSmith, ensure that you have a `LANGSMITH_API_KEY` set in your environment.\n",
    "\n",
    "Then set `os.environ['LANGCHAIN_TRACING_V2'] = 'false'` to `true`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['LANGCHAIN_TRACING_V2'] = 'false'\n",
    "os.environ['LANGCHAIN_PROJECT'] = 'Graphiti LangGraph Tutorial'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure Graphiti\n",
    "\n",
    "Ensure that you have `neo4j` running and a database created. Ensure that you've configured the following in your environment.\n",
    "\n",
    "```bash\n",
    "NEO4J_URI=\n",
    "NEO4J_USER=\n",
    "NEO4J_PASSWORD=\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Graphiti\n",
    "\n",
    "from graphiti_core import Graphiti\n",
    "from graphiti_core.edges import EntityEdge\n",
    "from graphiti_core.nodes import EpisodeType\n",
    "from graphiti_core.utils.bulk_utils import RawEpisode\n",
    "from graphiti_core.utils.maintenance.graph_data_operations import clear_data\n",
    "\n",
    "neo4j_uri = os.environ.get('NEO4J_URI', 'bolt://localhost:7687')\n",
    "neo4j_user = os.environ.get('NEO4J_USER', 'neo4j')\n",
    "neo4j_password = os.environ.get('NEO4J_PASSWORD', 'password')\n",
    "\n",
    "client = Graphiti(\n",
    "    neo4j_uri,\n",
    "    neo4j_user,\n",
    "    neo4j_password,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating a database schema \n",
    "\n",
    "The following is only required for the first run of this notebook or when you'd like to start your database over.\n",
    "\n",
    "**IMPORTANT**: `clear_data` is destructive and will wipe your entire database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: This will clear the database\n",
    "await clear_data(client.driver)\n",
    "await client.build_indices_and_constraints()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Shoe Data into the Graph\n",
    "\n",
    "Load several shoe and related products into the Graphiti. This may take a while.\n",
    "\n",
    "\n",
    "**IMPORTANT**: This only needs to be done once. If you run `clear_data` you'll need to rerun this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def ingest_products_data(client: Graphiti):\n",
    "    script_dir = Path.cwd().parent\n",
    "    json_file_path = script_dir / 'data' / 'manybirds_products.json'\n",
    "\n",
    "    with open(json_file_path) as file:\n",
    "        products = json.load(file)['products']\n",
    "\n",
    "    episodes: list[RawEpisode] = [\n",
    "        RawEpisode(\n",
    "            name=product.get('title', f'Product {i}'),\n",
    "            content=str({k: v for k, v in product.items() if k != 'images'}),\n",
    "            source_description='ManyBirds products',\n",
    "            source=EpisodeType.json,\n",
    "            reference_time=datetime.now(),\n",
    "        )\n",
    "        for i, product in enumerate(products)\n",
    "    ]\n",
    "\n",
    "    await client.add_episode_bulk(episodes)\n",
    "\n",
    "\n",
    "await ingest_products_data(client)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a user node in the Graphiti graph\n",
    "\n",
    "In your own app, this step could be done later once the user has identified themselves and made their sales intent known. We do this here so we can configure the agent with the user's `node_uuid`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_name = 'jess'\n",
    "\n",
    "await client.add_episode(\n",
    "    name='User Creation',\n",
    "    episode_body=(f'{user_name} is interested in buying a pair of shoes'),\n",
    "    source=EpisodeType.text,\n",
    "    reference_time=datetime.now(),\n",
    "    source_description='SalesBot',\n",
    ")\n",
    "\n",
    "# let's get Jess's node uuid\n",
    "nl = await client.get_nodes_by_query(user_name)\n",
    "\n",
    "user_node_uuid = nl[0].uuid\n",
    "\n",
    "# and the ManyBirds node uuid\n",
    "nl = await client.get_nodes_by_query('ManyBirds')\n",
    "manybirds_node_uuid = nl[0].uuid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def edges_to_facts_string(entities: list[EntityEdge]):\n",
    "    return '-' + '\\n- '.join([edge.fact for edge in entities])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import AIMessage, SystemMessage\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import END, START, StateGraph, add_messages\n",
    "from langgraph.prebuilt import ToolNode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `get_shoe_data` Tool\n",
    "\n",
    "The agent will use this to search the Graphiti graph for information about shoes. We center the search on the `manybirds_node_uuid` to ensure we rank shoe-related data over user data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "async def get_shoe_data(query: str) -> str:\n",
    "    \"\"\"Search the graphiti graph for information about shoes\"\"\"\n",
    "    edge_results = await client.search(\n",
    "        query,\n",
    "        center_node_uuid=manybirds_node_uuid,\n",
    "        num_results=10,\n",
    "    )\n",
    "    return edges_to_facts_string(edge_results)\n",
    "\n",
    "\n",
    "tools = [get_shoe_data]\n",
    "tool_node = ToolNode(tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model='gpt-4o-mini', temperature=0).bind_tools(tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [ToolMessage(content=\"-The product 'Men's SuperLight Wool Runners - Dark Grey (Medium Grey Sole)' is made of Wool.\\n- Women's Tree Breezers Knit - Rugged Beige (Hazy Beige Sole) has sizing options related to women's move shoes half sizes.\\n- TinyBirds Wool Runners - Little Kids - Natural Black (Blizzard Sole) is a type of Shoes.\\n- The product 'Men's SuperLight Wool Runners - Dark Grey (Medium Grey Sole)' belongs to the category Shoes.\\n- The product 'Men's SuperLight Wool Runners - Dark Grey (Medium Grey Sole)' uses SuperLight Foam technology.\\n- TinyBirds Wool Runners - Little Kids - Natural Black (Blizzard Sole) is sold by Manybirds.\\n- Jess is interested in buying a pair of shoes.\\n- TinyBirds Wool Runners - Little Kids - Natural Black (Blizzard Sole) has the handle TinyBirds-wool-runners-little-kids.\\n- ManyBirds Men's Couriers are a type of Shoes.\\n- Women's Tree Breezers Knit - Rugged Beige (Hazy Beige Sole) belongs to the Shoes category.\", name='get_shoe_data', tool_call_id='call_EPpOpD75rdq9jKRBUsfRnfxx')]}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test the tool node\n",
    "await tool_node.ainvoke({'messages': [await llm.ainvoke('wool shoes')]})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chatbot Function Explanation\n",
    "\n",
    "The chatbot uses Graphiti to provide context-aware responses in a shoe sales scenario. Here's how it works:\n",
    "\n",
    "1. **Context Retrieval**: It searches the Graphiti graph for relevant information based on the latest message, using the user's node as the center point. This ensures that user-related facts are ranked higher than other information in the graph.\n",
    "\n",
    "2. **System Message**: It constructs a system message incorporating facts from Graphiti, setting the context for the AI's response.\n",
    "\n",
    "3. **Knowledge Persistence**: After generating a response, it asynchronously adds the interaction to the Graphiti graph, allowing future queries to reference this conversation.\n",
    "\n",
    "This approach enables the chatbot to maintain context across interactions and provide personalized responses based on the user's history and preferences stored in the Graphiti graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class State(TypedDict):\n",
    "    messages: Annotated[list, add_messages]\n",
    "    user_name: str\n",
    "    user_node_uuid: str\n",
    "\n",
    "\n",
    "async def chatbot(state: State):\n",
    "    facts_string = None\n",
    "    if len(state['messages']) > 0:\n",
    "        last_message = state['messages'][-1]\n",
    "        graphiti_query = f'{\"SalesBot\" if isinstance(last_message, AIMessage) else state[\"user_name\"]}: {last_message.content}'\n",
    "        # search graphiti using Jess's node uuid as the center node\n",
    "        # graph edges (facts) further from the Jess node will be ranked lower\n",
    "        edge_results = await client.search(\n",
    "            graphiti_query, center_node_uuid=state['user_node_uuid'], num_results=5\n",
    "        )\n",
    "        facts_string = edges_to_facts_string(edge_results)\n",
    "\n",
    "    system_message = SystemMessage(\n",
    "        content=f\"\"\"You are a skillfull shoe salesperson working for ManyBirds. Review information about the user and their prior conversation below and respond accordingly.\n",
    "        Keep responses short and concise. And remember, always be selling (and helpful!)\n",
    "\n",
    "        Things you'll need to know about the user in order to close a sale:\n",
    "        - the user's shoe size\n",
    "        - any other shoe needs? maybe for wide feet?\n",
    "        - the user's preferred colors and styles\n",
    "        - their budget\n",
    "\n",
    "        Ensure that you ask the user for the above if you don't already know.\n",
    "\n",
    "        Facts about the user and their conversation:\n",
    "        {facts_string or 'No facts about the user and their conversation'}\"\"\"\n",
    "    )\n",
    "\n",
    "    messages = [system_message] + state['messages']\n",
    "\n",
    "    response = await llm.ainvoke(messages)\n",
    "\n",
    "    # add the response to the graphiti graph.\n",
    "    # this will allow us to use the graphiti search later in the conversation\n",
    "    # we're doing async here to avoid blocking the graph execution\n",
    "    asyncio.create_task(\n",
    "        client.add_episode(\n",
    "            name='Chatbot Response',\n",
    "            episode_body=f\"{state['user_name']}: {state['messages'][-1]}\\nSalesBot: {response.content}\",\n",
    "            source=EpisodeType.message,\n",
    "            reference_time=datetime.now(),\n",
    "            source_description='Chatbot',\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return {'messages': [response]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the Agent\n",
    "\n",
    "This section sets up the Agent's LangGraph graph:\n",
    "\n",
    "1. **Graph Structure**: It defines a graph with nodes for the agent (chatbot) and tools, connected in a loop.\n",
    "\n",
    "2. **Conditional Logic**: The `should_continue` function determines whether to end the graph execution or continue to the tools node based on the presence of tool calls.\n",
    "\n",
    "3. **Memory Management**: It uses a MemorySaver to maintain conversation state across turns. This is in addition to using Graphiti for facts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_builder = StateGraph(State)\n",
    "\n",
    "memory = MemorySaver()\n",
    "\n",
    "\n",
    "# Define the function that determines whether to continue or not\n",
    "async def should_continue(state, config):\n",
    "    messages = state['messages']\n",
    "    last_message = messages[-1]\n",
    "    # If there is no function call, then we finish\n",
    "    if not last_message.tool_calls:\n",
    "        return 'end'\n",
    "    # Otherwise if there is, we continue\n",
    "    else:\n",
    "        return 'continue'\n",
    "\n",
    "\n",
    "graph_builder.add_node('agent', chatbot)\n",
    "graph_builder.add_node('tools', tool_node)\n",
    "\n",
    "graph_builder.add_edge(START, 'agent')\n",
    "graph_builder.add_conditional_edges('agent', should_continue, {'continue': 'tools', 'end': END})\n",
    "graph_builder.add_edge('tools', 'agent')\n",
    "\n",
    "\n",
    "graph = graph_builder.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our LangGraph agent graph is illustrated below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "/9j/4AAQSkZJRgABAQAAAQABAAD/4gHYSUNDX1BST0ZJTEUAAQEAAAHIAAAAAAQwAABtbnRyUkdCIFhZWiAH4AABAAEAAAAAAABhY3NwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAA9tYAAQAAAADTLQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAlkZXNjAAAA8AAAACRyWFlaAAABFAAAABRnWFlaAAABKAAAABRiWFlaAAABPAAAABR3dHB0AAABUAAAABRyVFJDAAABZAAAAChnVFJDAAABZAAAAChiVFJDAAABZAAAAChjcHJ0AAABjAAAADxtbHVjAAAAAAAAAAEAAAAMZW5VUwAAAAgAAAAcAHMAUgBHAEJYWVogAAAAAAAAb6IAADj1AAADkFhZWiAAAAAAAABimQAAt4UAABjaWFlaIAAAAAAAACSgAAAPhAAAts9YWVogAAAAAAAA9tYAAQAAAADTLXBhcmEAAAAAAAQAAAACZmYAAPKnAAANWQAAE9AAAApbAAAAAAAAAABtbHVjAAAAAAAAAAEAAAAMZW5VUwAAACAAAAAcAEcAbwBvAGcAbABlACAASQBuAGMALgAgADIAMAAxADb/2wBDAAMCAgMCAgMDAwMEAwMEBQgFBQQEBQoHBwYIDAoMDAsKCwsNDhIQDQ4RDgsLEBYQERMUFRUVDA8XGBYUGBIUFRT/2wBDAQMEBAUEBQkFBQkUDQsNFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBT/wAARCADuAO4DASIAAhEBAxEB/8QAHQABAAICAwEBAAAAAAAAAAAAAAYHBQgCAwQBCf/EAFUQAAEDBAADAgYKChAGAgMAAAECAwQABQYRBxIhEzEIFBYiQVEVMlNVYXGSlNHTIzZWc3SBk5Wh1AkXJDM3OEJEVGKRsbKztME1UnJ1goMYJVeio//EABsBAQEAAwEBAQAAAAAAAAAAAAABAgMEBQYH/8QAMxEBAAECAgcHAwMFAQAAAAAAAAECEQMxBBIhQVFSoQUTFGFxkdEjscEVM1MyQmKB4fD/2gAMAwEAAhEDEQA/AP1TpSlApSlApSlApSlAr4pQSkkkADqSfRWMvt6NqbZajsGbcZSi3GihXKFq9Klq0eRtI6qVo6GgApRSlWNRhEe5KS/kDpvsnYV2Tw1FaI9DbOynW/SrmV/W7q3U0RbWrm0dVtxZNzJrO0spXdYKFDvSqSgEfprj5VWT34gfOkfTXxvErG0gIRZrehA7kpitgD9FcvJay+9ED5sj6Ky+j59DY+eVVk9+IHzpH008qrJ78QPnSPpr75LWX3ogfNkfRTyWsvvRA+bI+in0fPoux88qrJ78QPnSPpp5VWT34gfOkfTX3yWsvvRA+bI+inktZfeiB82R9FPo+fQ2PnlVZPfiB86R9Nd0a/2uY4G2LlEfWegS2+lRP4ga6vJay+9ED5sj6K6pOGY/MaLUixW19s96HIbah/YRT6Pn0NjM0qMeS72ODt8eedbaQNqtDzpVHdHqQVbLSvVykI9aT3jNWi7MXqA3KYC0BW0radTyuNLB0pCh6FA7BFYVUREa1M3j/wBmlntpSlakKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQRjHdXXJ7/c3NK8WeFsjd/mNoQlbnxFTiiDrvDafVoSeoxhw8UuGTQVbDjVyU+Nj2yHUIcCh6xsrT8aTWWyDI7Tidqeul8ukKzWxkpDs24SEMMo5iEp5lrIA2SANnqSBXRj/123Wj7LObI1j8gvsLF7Dcrzcnewt1ujOzJLvKVcjTaCtatDqdAE6FQ/8A+QvCv/8AJeH/AJ+i/WVwe404BksWVa7HluKZTd5TDjcWyM3uKtU5fIdM6CldFdx6HQJ6VzogmdeEzPj8BsszrHsLyCDIt8BqZAN+hNIZfbdBKHxyv+c2kDmUNhYBT5vWpxN4vTbbiVtu73DzMX5kx1TPsPEiRnpbXKN9o5yvlpKDroS5s9BrfSqStnA/OL7w64rYxFsL2C4zerKiLZMYud4bnojTtOF1TK0KWGWFfYwEb6HZCUjpUi4g2PPOJjGDz75w2lyrLbnZTd4wo3uJ+61qaa8XkqUHQ060hYeHZqVvqFcp1qgms/wnMUhYVi+TNwb3Nj5DdjY40GNC3MamgPczLrSlApUFsLQdb84j+TtQwrPhA5G/xztOInh9kEa1zLH4+6263E8ajuGUhrtnCJRT2KEk7CeZe1DQVUJwLglmVjsOC26Ri8e1JsvEaXfXY0Oay5HYt7rMpSFNnaSUoVIQ1y8oVtJPLy6NWRnNhyyw8ebJnFhxtWU216wO2GZHYmsxnYijJbeQ99lUkLRoKBCSVd3Q0FyUqAL8ILha2tSF8SsQStJ0UqvsUEH1fvlFeEHwtQopVxKxBKgdEG/Rdj/+lBP6jEfVo4gPxkAJZu8NUzkG/wB+ZU22tXq2pDrI/wDX8dSKLKZmxmpEd1D8d5AcbdaUFIWkjYUCOhBHXYqPSUeO8R7fyg8tvtj63Ty9AXnWg31+Jh3p8VdGD/dE5Wn/AJ1ssJNSlK50KUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQYG926RFuTV8tzPby2m+wlRgdGUwCVBKSenaJUSU76HmWk65+ZPvt10gZFDLsZxElnfKttaSFNq7+VaFAKQoelKgCPSK99Ya7Yha7xLEt1lcecAB45DeXHeIHcCtBBUB181Wx1PTqa3xVTVERibt6+rIexkM/zRj8mPorkiBGaWFIjtIUO5SUAEVHvId0dEZPfkJ9A8ZbV+lTZP6a+eRD/wB1N+/LtfVVe7w+fpK2jilNKi3kQ/8AdTfvy7X1VVn4Sk2+cJ+B2WZbZMouyrra46HWBKW0tvZdQk8yQgE9FH007vD5+klo4r1pUJseKS7jZLfLdym+9q/HbdXyvNAbUkE6+x/DXt8iH/upv35dr6qnd4fP0ktHFIPY2If5qz+TH0U9jYY/mrH5MfRUf8iH/upv35dr6qvvkMpY5XsjvzyPSnxsN7/GhKT+mmph8/SUtHFkrzkMWylqOB4xcHhqNb2CO1d9HQehI9KjpKfSRXHHbM7bW5MmYpDt0nOB+UtrZQFBISlCN9eRKQAO7Z2rQKjXbZcbtuPpcEGKGlua7R5alOOuerncUSpX4yaydY1VUxGrR78T0KUpWlClKUClKUClKUClKUClKUClKUClKUClKUClKUClKUCqL8OH+KnxD/A2v9Q1V6VRfhw/xU+If4G1/qGqC38U+1az/gbP+AVlaxWKfatZ/wADZ/wCsrQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKovw4f4qfEP8AA2v9Q1V6VRfhw/xU+If4G1/qGqC38U+1az/gbP8AgFZWsVin2rWf8DZ/wCsrQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKVgsiyRdqfZgwYqZ90fQpxDK3OzbQhJAK3F6VyjZAAAJJ7hoKKcMb5mGzqDZNejcl76uumjR6641tkesrZNq/Jv9kj4E/tc8XhmNtjFFiyzmkOFI81qcP35J9XPsOdT1Kl66Jr9LfZzMP6DY/nT31dVx4QHCq8eEJw3l4leY9mhpW63JjTmnnVuRXkHotIKNHaStJHqWe7vrPwtfGPeCzX/APYsOCz1rtN+4nzkLaNxQqz21J2AthK0rec9RBcQhIPoLa63+qrcHtV+4e4fZsas1ssbFstUVuIwgyndlKRrmUez6qPUk+kkn01m/ZzMP6DY/nT31dPC18Y94LJvSoSm+5ekgm32RwD+SJjyd/j7I6/sNSLH783f4jiw0uNJYX2MmM51Uy4ADrY6EEEEEd4IPwVrrwK8ONacvKblmUpSlc6FKUoFKUoFKUoFKUoFKUoFKUoFKUoFKUoFKUoFKUoINNO+Js8egWeLr4NvSN/3CstWImfwnXD/ALPE/wA+TUA48ZZcLI3j9qseQX
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with suppress(Exception):\n",
    "    display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the Agent\n",
    "\n",
    "Let's test the agent with a single call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='What sizes do the TinyBirds Wool Runners in Natural Black come in?', id='6a940637-70a0-4c95-a4d7-4c4846909747'),\n",
       "  AIMessage(content='The TinyBirds Wool Runners in Natural Black are available in the following sizes for little kids: 5T, 6T, 8T, 9T, and 10T. \\n\\nDo you have a specific size in mind, or are you looking for something else? Let me know your needs, and I can help you find the perfect pair!', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 76, 'prompt_tokens': 314, 'total_tokens': 390}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_f33667828e', 'finish_reason': 'stop', 'logprobs': None}, id='run-d2f79c7f-4d41-4896-88dc-476a8e38bea8-0', usage_metadata={'input_tokens': 314, 'output_tokens': 76, 'total_tokens': 390})],\n",
       " 'user_name': 'jess',\n",
       " 'user_node_uuid': '186a845eee4849619d1e625b178d1845'}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await graph.ainvoke(\n",
    "    {\n",
    "        'messages': [\n",
    "            {\n",
    "                'role': 'user',\n",
    "                'content': 'What sizes do the TinyBirds Wool Runners in Natural Black come in?',\n",
    "            }\n",
    "        ],\n",
    "        'user_name': user_name,\n",
    "        'user_node_uuid': user_node_uuid,\n",
    "    },\n",
    "    config={'configurable': {'thread_id': uuid.uuid4().hex}},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Viewing the Graph\n",
    "\n",
    "At this stage, the graph would look something like this. The `jess` node is `INTERESTED_IN` the `TinyBirds Wool Runner` node. The image below was generated using Neo4j Desktop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABp4AAAPQCAYAAAAvigi5AAAAAXNSR0IArs4c6QAAAJZlWElmTU0AKgAAAAgABAEaAAUAAAABAAAAPgEbAAUAAAABAAAARgEoAAMAAAABAAIAAIdpAAQAAAABAAAATgAAAAAAAACQAAAAAQAAAJAAAAABAASShgAHAAAAEgAAAISgAQADAAAAAQABAACgAgAEAAAAAQAABp6gAwAEAAAAAQAAA9AAAAAAQVNDSUkAAABTY3JlZW5zaG90XTAm7gAAAAlwSFlzAAAWJQAAFiUBSVIk8AAAAddpVFh0WE1MOmNvbS5hZG9iZS54bXAAAAAAADx4OnhtcG1ldGEgeG1sbnM6eD0iYWRvYmU6bnM6bWV0YS8iIHg6eG1wdGs9IlhNUCBDb3JlIDYuMC4wIj4KICAgPHJkZjpSREYgeG1sbnM6cmRmPSJodHRwOi8vd3d3LnczLm9yZy8xOTk5LzAyLzIyLXJkZi1zeW50YXgtbnMjIj4KICAgICAgPHJkZjpEZXNjcmlwdGlvbiByZGY6YWJvdXQ9IiIKICAgICAgICAgICAgeG1sbnM6ZXhpZj0iaHR0cDovL25zLmFkb2JlLmNvbS9leGlmLzEuMC8iPgogICAgICAgICA8ZXhpZjpQaXhlbFlEaW1lbnNpb24+OTc2PC9leGlmOlBpeGVsWURpbWVuc2lvbj4KICAgICAgICAgPGV4aWY6UGl4ZWxYRGltZW5zaW9uPjE2OTQ8L2V4aWY6UGl4ZWxYRGltZW5zaW9uPgogICAgICAgICA8ZXhpZjpVc2VyQ29tbWVudD5TY3JlZW5zaG90PC9leGlmOlVzZXJDb21tZW50PgogICAgICA8L3JkZjpEZXNjcmlwdGlvbj4KICAgPC9yZGY6UkRGPgo8L3g6eG1wbWV0YT4Ka8cvqQAAABxpRE9UAAAAAgAAAAAAAAHoAAAAKAAAAegAAAHoAAJFvAqvVXMAAEAASURBVHgB7J0HYBXFFob/e9MhDdKoIRBq6CgIAhYElKaIiooCUgUVFaR3kd5RioiAhaIC0lG6CEjvHZLQW0gogRDSbt7M8nZJSAIpN7f+qyGzs7MzZ77Z3ezOmXOOLi4xORncSIAESIAESIAESIAESIAESCCXCdyIvIktW3elaSWkbEmElCuVJj+7GbIdP9/8aU5P2f7bbzZKc5wZ9kvg1OlwHDl2Cs7OTqhauTyKFilolTBOng7D0WOnUaVyCEqWKGaVfchNoeXzRz4HbOX+P37iDI6fDFWedy/WfS430bFuEiABGycQ/3kZINmg9FJfqBT0voEW1WOdez4kP7gHJCYYVy5HF+jyegOGBCTfu60xMEYjyXH3kXRqh1aVY9fZ0Ie8oO2bK3H82DEcP34MYaGhCA8XP2FhiImJgV6vR+HCRRBcsiRKlAhG+46dlDxzyWnP7X7wXksUCwrC8JGj0x2Djzu1x4svvYxWH7S2aEw6Kp4senwoHAmQAAmQAAmQAAmQAAnYDAF10je9DplqIlidqDW2siu9PjHPOggcOnICZ0LPwcM9L6pWKQ9/Px/rEPwxKWPux2L9xq3wyZ8PdWtXf+wodyUB9f6XSpr0lNPWSEl9rvKZZo2jR5lJwHIIxPeoACTEKQLpC5SE3p+LF3I6OlJRlnT60YIrx26/QF+6Vk6rzfH5fXr2EH8PjyM4uKT4CYaPry++nz4NazdsgptbnhzXzwpyRuDYsaPo0qkDlixbCX9//3QrmzN7Fk6eOIGx4yeme9xSMql4spSRoBwkQAIkQAIkQAIkQAIkYOMEFi/9K8MemnLSVJ18NmWbGXacB8xKYNeeg7h46Sp8ffIplk5eXh5mlScnje87cBRnz13E8zWroVDBgJxUZbPnqlaPtnbvU/lks5csO0YCJiMQ368GcO+W0p7ePwj6AsEma9tWG0q+fwdJoXu17jl9tRi6oMravrkSCQkJcHJy0ppPFs7Q6r/8Aub+PA+Bxahw1MCYKbF500b8+MNMzP/tD02CMaNG4o3mzVG2XIiS9/WQQdDpdBg8dJhWxhITVDxZ4qhQJhIgARIgARIgARIgARKwMQLqhO+TumUqqycpg6oEM2WbT+o7j5mWQFxcPHbuPqC4XStUKADVhHs9V1cX0wphxNYibkTh3227EVSsCJ6tVtGINdteVfLel9ZOtuSaLuXz1ZasuWzv6mOPSMByCSR80xDJEWcVAfU+haEvXNZyhbUSyZKjI5F07pAmrdPg9dD5BWn7lpTo2L4tPmzdFi+9XM+SxLJLWfbu2Y1vvh6KpStWaW72Fv3+G+bP+wVzhHLQ29sbrd59G7379ke1Z57F4UOHUL5CBTg4OFgcLyqeLG5IKBAJkAAJkAAJkAAJkAAJ2B4BdUX+k3pmSisE1erJ1iagn8SXxx4SuBN9VyidDuLu3XsoHlRUca+nF6tGrXmTSqc7d+4qLva8vT2tuSu5Lrv6LLI1pXNK5ZOt9S3XLwo2QAIkgIQprZAcukchofP0g0NQpSxR0eXxgr5YRejcRYzN+FgYoi7BcOlEluqwtcKGqMswXD6pdct53AHA1V3bt6TEqBHfCLduAejQqbMliWWXshgMBrR67x20a9cBrzZqrDEY8c0wXLhwDlOnz8T9+/fh5eWlHPvsky6oXbsO3v/gQ62spSSoeLKUkaAcJEACJEACJEACJEACJGCjBFJOiD6ti+ZQPpmyzaf1n8dzl8D1iEhF6STdzJQtE4wKIaVzt0ET1C7d60k3eyHlSkFey9yeTEBVOtuiZZDaNyrUn3wN8CgJkEBaAok/fwXD3hXKAZ2bOxxKPZe2UAY5jtUawfHZJoCYMJfu5XQevhB+wGCIvICE1d8hOfZuBmfadrbhWhgMEecedjKPJ5zH7LPYDv/+2wIcOngQI0ePxe3btxEeForw8HC89fY7iks3ixXcRgXbvWsnBvbvi66fdsOrrzVCXNwDLJg3DwsXzMOESVPwXM2HscLkWDVt1BDjJ05GzVrPWxwNKp4sbkgoEAmQAAmQAAmQAAmQAAnYFgHVwuBJvZITpepmShdYqss9W5yEVnny90MCFy5ewe69D13eVK5UDqWCg6weTWJiItZt3AZnZye8ULuG8tvqO5XLHVCVM7aqcLb1/uXy5cHqScBuCSStnoKkv6c+7L/eAY4VXsoUC4fiVeD0ahckHlqPxD1CcZWYAJ1rXuiDn4VT3fdhuHgc8au/zVRdtlYo6fwRJN+JULqlC6wIp15/WlQXHzx4gLNCuRQeHort27Zh184dyJMnD6KiohQ5fX398NOv85E//6N3dIvqgI0Lc+zYUcz5cRaOHjkirPSj4e7hgWHfjEilYFq9agWmfjsFq/5aR1d7Nn49sHskQAIkQAIkQAIkQAIkQALpEJAToXLz8/PRjkpllCWsyk9pjUX3VNrw2Fzi9JmzOHz0pAim7YiqIp5TYNFCNtHHo8dO4eTpcCWuk4zvxC1zBKTC2VYVT5KAquy35T5mbqRZigRIILMEDPtWIfGn7lpxhzK1oHPJo+1nlHB6qbWwjqqBB7PFuYbEVMWcm36hxIp6MEccS3iQ6ljmd4QrXBm7Jil13WnOd3B8ehl5kk4PJBvSnJ4bGUmndiA57r5Stb56czi2GZcbzWSrzskTJ2Dxot+FkZoB7u7uKBEcjBIlghFcsiSCg0uK/ZLw9KTr3mzBNfJJo0eNENZoBzB2/EQULRqYqvZeX3WHj68v+vYbkCrfUnZo8WQpI0E5SIAESIAESIAESIAESMCOCKgTo5ag7KGFgG1feIePnMTp0LNiYiWPonQK8BcugGxgkzGd1m/ahkIFA/B8zWo20CPTdUG1dLSE509u9Vp9xtKaM7cIs14SsC0CyVfPIGHko3gyDoEVoPMOeGonnV5uC4cyNRE3fxCS70amKi9d7un9ApU4R1IB41TnXcj4UfFr/m9Z9f/STg1FXCERFyrhn1+hc/OAc4t+SNy1FHohgxI3ytkVyTevIHHfaiSFizhJKTZ53KnGG9DlLyRc+t1T4kolbPtNqU8Wc3qpjVBIJSDx8AbR/vv
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/png": {
       "width": 850
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Image(filename='tinybirds-jess.png', width=850))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the Agent interactively\n",
    "\n",
    "The following code will run the agent in an event loop. Just enter a message into the box and click submit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation_output = widgets.Output()\n",
    "config = {'configurable': {'thread_id': uuid.uuid4().hex}}\n",
    "user_state = {'user_name': user_name, 'user_node_uuid': user_node_uuid}\n",
    "\n",
    "\n",
    "async def process_input(user_state: State, user_input: str):\n",
    "    conversation_output.append_stdout(f'\\nUser: {user_input}\\n')\n",
    "    conversation_output.append_stdout('\\nAssistant: ')\n",
    "\n",
    "    graph_state = {\n",
    "        'messages': [{'role': 'user', 'content': user_input}],\n",
    "        'user_name': user_state['user_name'],\n",
    "        'user_node_uuid': user_state['user_node_uuid'],\n",
    "    }\n",
    "\n",
    "    try:\n",
    "        async for event in graph.astream(\n",
    "            graph_state,\n",
    "            config=config,\n",
    "        ):\n",
    "            for value in event.values():\n",
    "                if 'messages' in value:\n",
    "                    last_message = value['messages'][-1]\n",
    "                    if isinstance(last_message, AIMessage) and isinstance(\n",
    "                        last_message.content, str\n",
    "                    ):\n",
    "                        conversation_output.append_stdout(last_message.content)\n",
    "    except Exception as e:\n",
    "        conversation_output.append_stdout(f'Error: {e}')\n",
    "\n",
    "\n",
    "def on_submit(b):\n",
    "    user_input = input_box.value\n",
    "    input_box.value = ''\n",
    "    asyncio.create_task(process_input(user_state, user_input))\n",
    "\n",
    "\n",
    "input_box = widgets.Text(placeholder='Type your message here...')\n",
    "submit_button = widgets.Button(description='Send')\n",
    "submit_button.on_click(on_submit)\n",
    "\n",
    "conversation_output.append_stdout('Asssistant: Hello, how can I help you find shoes today?')\n",
    "\n",
    "display(widgets.VBox([input_box, submit_button, conversation_output]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}