11.5-Smart_Agent-LangGraph.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6423f8f3-a592-4ee7-9969-39e38933be52",
   "metadata": {},
   "source": [
    "# Putting it all together using LangGraph (experimental)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d7fa9dc-64cb-4ee2-ae98-8cdb72293cbe",
   "metadata": {},
   "source": [
    "**Last notebook we created a brain langchain agent that acts as a supervisor of the experts (other agents with tools).** This concept of multi-agent architectures is a current field of high research interest.\n",
    "\n",
    "**Graphs are important in multi-agent systems** as they efficiently represent the interactions and relationships between different agents:\n",
    "- **Nodes (or Vertices):** Each agent, can perform specific tasks or make decisions.\n",
    "- **Edges:** Signify communication paths or decision flows between agents.\n",
    "\n",
    "This structure enables the division of complex problems into smaller, manageable tasks, where each agent can focus on a particular aspect. The advantages of this approach include:\n",
    "- Improved specialization and parallelization.\n",
    "- More robust and scalable solutions.\n",
    "\n",
    "**In our context, we can use the graph-based architecture** to build systems where agents coordinate through a defined protocol of interactions, enhancing both performance and flexibility. For instance, if one agent fails or generates suboptimal results, other agents in the system can take over or adjust their strategies accordingly, thus ensuring continuity and efficiency.\n",
    "\n",
    "[**LangGraph**](https://python.langchain.com/docs/langgraph/), specifically designed for building multi-agent systems with LLMs, leverages these graph concepts by allowing developers to define states and transitions explicitly. It treats each agent as a state machine, with transitions defined by the possible actions an agent can take in response to its environment or the inputs it receives. LangGraph provides tools for building these systems, where each node in the graph can be an agent capable of making decisions or performing actions based on the shared state of the system. The library supports creating complex workflows where agents interact through a shared state, making it easier to develop, test, and maintain each component in isolation before integrating them into the broader system.\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "30b81551-92ac-4f08-9c00-ba11981c67c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import operator\n",
    "from typing import TypedDict, Annotated, Sequence, Union, List\n",
    "from langchain_openai import AzureChatOpenAI\n",
    "from langchain_core.messages import HumanMessage, AIMessage, BaseMessage\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain.callbacks.manager import CallbackManager\n",
    "\n",
    "from langgraph.prebuilt import ToolExecutor, ToolNode\n",
    "from langgraph.graph import StateGraph, MessageGraph, END\n",
    "\n",
    "\n",
    "#custom libraries that we will use later in the app\n",
    "from common.utils import (\n",
    "    DocSearchAgent, \n",
    "    CSVTabularAgent, \n",
    "    SQLSearchAgent, \n",
    "    ChatGPTTool, \n",
    "    BingSearchAgent,\n",
    ")\n",
    "from common.callbacks import StdOutCallbackHandler\n",
    "from common.prompts import CUSTOM_CHATBOT_PROMPT, CUSTOM_CHATBOT_PREFIX\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv(\"credentials.env\")\n",
    "\n",
    "from IPython.display import Markdown, HTML, display \n",
    "\n",
    "def printmd(string):\n",
    "    display(Markdown(string))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "67cd1e3e-8527-4a8f-ba90-e700ae7b20ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"OPENAI_API_VERSION\"] = os.environ[\"AZURE_OPENAI_API_VERSION\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d374f5a-620e-4bb5-abbb-3edb49d82c9b",
   "metadata": {},
   "source": [
    "### Define our Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "643d1650-6416-46fd-8b21-f5fb298ec063",
   "metadata": {},
   "outputs": [],
   "source": [
    "cb_handler = StdOutCallbackHandler()\n",
    "cb_manager = CallbackManager(handlers=[cb_handler])\n",
    "\n",
    "COMPLETION_TOKENS = 2000\n",
    "\n",
    "llm = AzureChatOpenAI(deployment_name=os.environ[\"GPT4o_DEPLOYMENT_NAME\"], \n",
    "                      temperature=0.5, max_tokens=COMPLETION_TOKENS)\n",
    "\n",
    "# Uncomment below if you want to see the answers streaming\n",
    "# llm = AzureChatOpenAI(deployment_name=os.environ[\"GPT35_DEPLOYMENT_NAME\"], temperature=0.5, \n",
    "#                        max_tokens=COMPLETION_TOKENS, streaming=True, callback_manager=cb_manager)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56b56a94-0471-41c3-b441-3a73ff5dedfc",
   "metadata": {},
   "source": [
    "### Define the Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a6a4cc93-2dd6-45eb-ac5b-5af2d31809dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_indexes = [\"srch-index-files\", \"srch-index-csv\"]\n",
    "doc_search = DocSearchAgent(llm=llm, indexes=doc_indexes,\n",
    "                           k=6, reranker_th=1,\n",
    "                           sas_token=os.environ['BLOB_SAS_TOKEN'],\n",
    "                           name=\"docsearch\",\n",
    "                           description=\"useful when the questions includes the term: docsearch\",\n",
    "                           verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "08e3bfa4-98c0-4b6f-a918-720f50a2f484",
   "metadata": {},
   "outputs": [],
   "source": [
    "book_indexes = [\"srch-index-books\"]\n",
    "book_search = DocSearchAgent(llm=llm, indexes=book_indexes,\n",
    "                           k=10, reranker_th=1,\n",
    "                           sas_token=os.environ['BLOB_SAS_TOKEN'],\n",
    "                           name=\"booksearch\",\n",
    "                           description=\"useful when the questions includes the term: booksearch\",\n",
    "                           verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0f0ae466-aff8-4cdf-80d3-ef2c61867fc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# BingSearchAgent is a langchain Tool class to use the Bing Search API (https://www.microsoft.com/en-us/bing/apis/bing-web-search-api)\n",
    "www_search = BingSearchAgent(llm=llm, k=5,\n",
    "                             name=\"bing\",\n",
    "                             description=\"useful when the questions includes the term: bing\",\n",
    "                             verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "78edb304-c4a2-4f10-8ded-936e9141aa02",
   "metadata": {},
   "outputs": [],
   "source": [
    "## CSVTabularAgent is a custom Tool class crated to Q&A over CSV files\n",
    "file_url = \"./data/all-states-history.csv\"\n",
    "csv_search = CSVTabularAgent(path=file_url, llm=llm,\n",
    "                             name=\"csvfile\",\n",
    "                             description=\"useful when the questions includes the term: csvfile\",\n",
    "                             verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b9d54cc5-41bc-43c3-a91d-12fc3a2446ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "## SQLDbAgent is a custom Tool class created to Q&A over a MS SQL Database\n",
    "sql_search = SQLSearchAgent(llm=llm, k=30,\n",
    "                            name=\"sqlsearch\",\n",
    "                            description=\"useful when the questions includes the term: sqlsearch\",\n",
    "                            verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "65465173-92f6-489d-9b48-58d109c5723e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## ChatGPTTool is a custom Tool class created to talk to ChatGPT knowledge\n",
    "chatgpt_search = ChatGPTTool(llm=llm,\n",
    "                             name=\"chatgpt\",\n",
    "                            description=\"useful when the questions includes the term: chatgpt\",\n",
    "                            verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d018c884-5c91-4a35-90e3-6a5a6e510c25",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "tools = [www_search, sql_search, doc_search, csv_search, chatgpt_search, book_search]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a14154d-a262-418c-bc26-27e06822d160",
   "metadata": {},
   "source": [
    "### Bind the Tools to the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4580ffef-366d-4e6e-b400-500095d7b5de",
   "metadata": {},
   "source": [
    "We should make sure the model knows that it has these tools available to call. We can do this by converting the LangChain tools into the format for OpenAI tool calling, and then bind them to the model class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8d9101ab-3a56-4e1e-be35-33d0d3ba12ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_with_tools = llm.bind_tools(tools) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2083506-f25c-492c-869d-fb97f66baa34",
   "metadata": {},
   "source": [
    "Let's test our model with a few messages to see if effectively calls the function sqlsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "89e53c87-077d-4cb3-b631-084166afc6e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_qj9So05P8CiSAj6XxaW7iz6c', 'function': {'arguments': '{\"query\":\"how many deaths in the east coast\"}', 'name': 'sqlsearch'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 20, 'prompt_tokens': 430, 'total_tokens': 450}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'tool_calls', 'logprobs': None, 'content_filter_results': {}}, id='run-cf36f0cd-b62c-4081-878e-4c354e80139e-0', tool_calls=[{'name': 'sqlsearch', 'args': {'query': 'how many deaths in the east coast'}, 'id': 'call_qj9So05P8CiSAj6XxaW7iz6c'}], usage_metadata={'input_tokens': 430, 'output_tokens': 20, 'total_tokens': 450})"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm_with_tools.invoke([HumanMessage(content=\"what is your name\"), \n",
    "              AIMessage(content='My name is Assistant. How can I assist you today?'),\n",
    "              HumanMessage(content=\"sqlsearch, how many deaths in the east coast?\")])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7301e40f-42f6-4f08-a65d-845ad2dadda8",
   "metadata": {},
   "source": [
    "This looks correct!,  the LLM responded with a AIMessage that has no content, but instead has a **tool_call**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a055e9f-8060-4ae5-94ae-3d966355392f",
   "metadata": {},
   "source": [
    "### Define agent state\n",
    "\n",
    "A graph is parameterized by a state object that it passes around to each node. Each node then returns operations to update that state. These operations can either SET specific attributes on the state (e.g. overwrite the existing values) or ADD to the existing attribute. Whether to set or add is denoted by annotating the state object you construct the graph with.\n",
    "\n",
    "For our case, the state we will track will just be a list of messages. We want each node to just add messages to that list. Therefore, we will use a TypedDict with one key (messages) and annotate it so that the messages attribute is always added to with the second parameter (operator.add)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ae34e42a-2522-4b3a-9a75-559c5d72fd0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AgentState(TypedDict):\n",
    "    messages: Annotated[Sequence[BaseMessage], operator.add]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08f2ef1c-fd0a-43f0-82bb-1ce466812e44",
   "metadata": {},
   "source": [
    "### Define the nodes\n",
    "\n",
    "We now need to define a few different nodes in our graph. In langgraph, **a node can be either a function or a runnable**. There are two main nodes we need for this:\n",
    "\n",
    "**The agent**: responsible for deciding what (if any) actions to take.\n",
    "A function to invoke tools: if the agent decides to take an action, this node will then execute that action.\n",
    "We will also need to define some edges. Some of these edges may be conditional. The reason they are conditional is that based on the output of a node, one of several paths may be taken. The path that is taken is not known until that node is run (the LLM decides).\n",
    "\n",
    "**Conditional Edge**: after the agent is called, we should either:\n",
    "\n",
    "    a. If the agent said to take an action, then the function to invoke tools should be called\n",
    "\n",
    "    b. If the agent said that it was finished, then it should finish\n",
    "\n",
    "**Normal Edge**: after the tools are invoked, it should always go back to the agent to decide what to do next\n",
    "\n",
    "Let's define the nodes, as well as a function to decide how what conditional edge to take."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "41063cdf-f337-44d4-b48a-46ae4170c987",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the function that determines whether to continue or not\n",
    "def should_continue(state):\n",
    "    \n",
    "    messages = state[\"messages\"]\n",
    "    \n",
    "    last_message = messages[-1]\n",
    "    # If there are no tool calls, then we finish\n",
    "    if not last_message.tool_calls:\n",
    "        return \"end\"\n",
    "   # Otherwise if there is, we check if it's suppose to return direct the result of the tool or not (based on the prompt)\n",
    "    else:\n",
    "        arguments = last_message.tool_calls[0][\"args\"]\n",
    "        if arguments.get(\"return_direct\", False):\n",
    "            return \"final\"\n",
    "        else:\n",
    "            return \"continue\"\n",
    "\n",
    "    \n",
    "\n",
    "# Define the function that calls the supervisor chain\n",
    "async def supervisor_node(state):\n",
    "    \n",
    "    messages = state[\"messages\"]\n",
    "    \n",
    "    PROMPT = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\"system\", CUSTOM_CHATBOT_PREFIX),\n",
    "            (\"human\", \"{question}\")\n",
    "        ]\n",
    "    )\n",
    "    \n",
    "    chain = (\n",
    "        {\n",
    "            \"question\": lambda x: x[\"question\"],\n",
    "        }\n",
    "        | PROMPT\n",
    "        | llm_with_tools\n",
    "    )\n",
    "    \n",
    "    response = await chain.ainvoke({\"question\": messages})\n",
    "    \n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "# Define the function to execute tools\n",
    "tool_node = ToolNode(tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f17a00e-0d0c-4166-b3c8-37bc22556664",
   "metadata": {},
   "source": [
    "### Define the Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a8f23fdd-7a8d-4f2c-900a-d11b20fe09cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a new graph\n",
    "workflow = StateGraph(AgentState)\n",
    "\n",
    "\n",
    "# Define the two nodes we will cycle between\n",
    "workflow.add_node(\"supervisor\", supervisor_node)\n",
    "workflow.add_node(\"tools\", tool_node)\n",
    "# We add a separate node for any tool call where return_direct=True. The reason this is needed is that after this node we want to end, \n",
    "# while after other tool calls we want to go back to the LLM.\n",
    "workflow.add_node(\"tools_final\", tool_node)\n",
    "\n",
    "# Set the entrypoint as `supervisor`\n",
    "# This means that this node is the first one called\n",
    "workflow.set_entry_point(\"supervisor\")\n",
    "\n",
    "# We now add a conditional edge\n",
    "workflow.add_conditional_edges(\n",
    "    # First, we define the start node. We use `supervisor`.\n",
    "    # This means these are the edges taken after the `supervisor` node is called.\n",
    "    \"supervisor\",\n",
    "    # Next, we pass in the function that will determine which node is called next.\n",
    "    should_continue,\n",
    "    # Finally we pass in a mapping.\n",
    "    # The keys are strings, and the values are other nodes.\n",
    "    # END is a special node marking that the graph should finish.\n",
    "    # What will happen is we will call `should_continue`, and then the output of that\n",
    "    # will be matched against the keys in this mapping.\n",
    "    # Based on which one it matches, that node will then be called.\n",
    "    {\n",
    "        # If `tools`, then we call the tool node.\n",
    "        \"continue\": \"tools\",\n",
    "        # Final call\n",
    "        \"final\": \"tools_final\",\n",
    "        # Otherwise we finish.\n",
    "        \"end\": END,\n",
    "    },\n",
    ")\n",
    "\n",
    "# We now add a normal edge from `tools` to `supervisor`.\n",
    "# This means that after `tools` is called, `supervisor` node is called next.\n",
    "workflow.add_edge(\"tools\", \"supervisor\")\n",
    "# and from `tools_final` to END\n",
    "workflow.add_edge(\"tools_final\", END)\n",
    "\n",
    "# Finally, we compile it!\n",
    "# This compiles it into a LangChain Runnable,\n",
    "# meaning you can use it as you would any other runnable\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1338f71-a718-4d12-b461-8083b8cf3b76",
   "metadata": {},
   "source": [
    "### Define a utility print function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1114992c-d50b-479d-8ee5-99c4f400de23",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def print_events(query):\n",
    "    inputs = {\"messages\": [HumanMessage(content=query)]}\n",
    "    event_names = [\"LangGraph\", \"supervisor\", \"call_supervisor_chain\", \"should_continue\", \"tools\", \"tools_final\"]\n",
    "\n",
    "    async for event in app.astream_events(inputs, version=\"v1\"):\n",
    "        kind = event[\"event\"]\n",
    "        # print(event)\n",
    "        if kind == \"on_chain_start\":\n",
    "            if event[\"name\"] in event_names:\n",
    "                print(\"\\n=======================\\n\")\n",
    "                print(f\"Starting: {event['name']}\")\n",
    "                print(event)\n",
    "\n",
    "        elif kind == \"on_chain_end\":\n",
    "            if event[\"name\"] == \"LangGraph\":  \n",
    "                print(\"\\n=======================\\n\")\n",
    "                # print(event)\n",
    "                try:\n",
    "                    printmd(event['data']['output']['supervisor']['messages'][0].content)\n",
    "                except:\n",
    "                    print(event['data']['output'])\n",
    "                    printmd(event['data']['output'][1]['tools_final']['messages'][0].content)\n",
    "        if kind == \"on_chat_model_stream\":\n",
    "            content = event[\"data\"][\"chunk\"].content\n",
    "            # Empty content in the context of OpenAI means that the model is asking for a tool to be invoked.\n",
    "            # So we only print non-empty content\n",
    "            if content:\n",
    "                print(content, end=\"|\")\n",
    "        elif kind == \"on_tool_start\":\n",
    "            print(\"\\n=======================\\n\")\n",
    "            print(f\"Starting tool: {event['name']}\")\n",
    "            print(event)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d685b6-6b58-4065-a0b2-8864286324bb",
   "metadata": {},
   "source": [
    "### Use the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "df0f3c85-1bfa-49f5-bd70-5f72f6bd8ecf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/anaconda/envs/azureml_py310_sdkv2/lib/python3.10/site-packages/langchain_core/_api/beta_decorator.py:87: LangChainBetaWarning: This API is in beta and may change in the future.\n",
      "  warn_beta(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=======================\n",
      "\n",
      "Starting: LangGraph\n",
      "{'event': 'on_chain_start', 'run_id': '27923040-d0ca-4874-b51b-9bff0ce8ddec', 'name': 'LangGraph', 'tags': [], 'metadata': {}, 'data': {'input': {'messages': [HumanMessage(content='Hello there, how are you? My name is Pablo Marin')]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: supervisor\n",
      "{'event': 'on_chain_start', 'name': 'supervisor', 'run_id': 'c85f6126-5e16-4987-93e7-8ab6e1cab6d9', 'tags': ['graph:step:1'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: should_continue\n",
      "{'event': 'on_chain_start', 'name': 'should_continue', 'run_id': 'd6ade7d2-06a3-4203-bdc7-67a15ffa846d', 'tags': ['seq:step:3'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {'input': {'messages': [HumanMessage(content='Hello there, how are you? My name is Pablo Marin'), AIMessage(content=\"Hello Pablo Marin! I'm Jarvis, your assistant. I'm here to help you with any questions or tasks you have. How can I assist you today?\", response_metadata={'token_usage': {'completion_tokens': 32, 'prompt_tokens': 875, 'total_tokens': 907}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'stop', 'logprobs': None, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}, id='run-fc6dfcea-bd98-4c97-bfab-d2ecdcc3b80b-0', usage_metadata={'input_tokens': 875, 'output_tokens': 32, 'total_tokens': 907})]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "Hello Pablo Marin! I'm Jarvis, your assistant. I'm here to help you with any questions or tasks you have. How can I assist you today?"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "await print_events(\"Hello there, how are you? My name is Pablo Marin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f66b2db6-ea88-4473-8f5f-6ec0db664ced",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=======================\n",
      "\n",
      "Starting: LangGraph\n",
      "{'event': 'on_chain_start', 'run_id': '429c72a4-0f32-4be7-b70b-95e2e883acc9', 'name': 'LangGraph', 'tags': [], 'metadata': {}, 'data': {'input': {'messages': [HumanMessage(content=\"bing, What's the oldest parrot alive, and how much longer is that than the average lifespan of a parrot?\")]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: supervisor\n",
      "{'event': 'on_chain_start', 'name': 'supervisor', 'run_id': 'f11c02d8-e966-4117-827a-13de54c7e36b', 'tags': ['graph:step:1'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: should_continue\n",
      "{'event': 'on_chain_start', 'name': 'should_continue', 'run_id': '0e1e34e5-5bf4-4be7-aeb3-652bd6431ab2', 'tags': ['seq:step:3'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {'input': {'messages': [HumanMessage(content=\"bing, What's the oldest parrot alive, and how much longer is that than the average lifespan of a parrot?\"), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_HXBAx6MHIpBrWoVaTZcCgWZw', 'function': {'arguments': '{\"query\":\"oldest parrot alive\"}', 'name': 'bing'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 18, 'prompt_tokens': 887, 'total_tokens': 905}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'tool_calls', 'logprobs': None, 'content_filter_results': {}}, id='run-a007cfce-535f-49b9-853d-d04a3e6abfe2-0', tool_calls=[{'name': 'bing', 'args': {'query': 'oldest parrot alive'}, 'id': 'call_HXBAx6MHIpBrWoVaTZcCgWZw'}], usage_metadata={'input_tokens': 887, 'output_tokens': 18, 'total_tokens': 905})]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: tools\n",
      "{'event': 'on_chain_start', 'name': 'tools', 'run_id': '7f66d307-1b92-48e2-b557-aa749866b8e1', 'tags': ['graph:step:2'], 'metadata': {'langgraph_step': 2, 'langgraph_node': 'tools', 'langgraph_triggers': ['branch:supervisor:should_continue:tools'], 'langgraph_task_idx': 0}, 'data': {}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting tool: bing\n",
      "{'event': 'on_tool_start', 'name': 'bing', 'run_id': '0b11499a-75dc-469f-9b8e-df2908a2c46f', 'tags': ['seq:step:1'], 'metadata': {'langgraph_step': 2, 'langgraph_node': 'tools', 'langgraph_triggers': ['branch:supervisor:should_continue:tools'], 'langgraph_task_idx': 0}, 'data': {'input': {'query': 'oldest parrot alive'}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: supervisor\n",
      "{'event': 'on_chain_start', 'name': 'supervisor', 'run_id': 'a7d8af8e-f9c5-4e37-9b14-9a68cef2887c', 'tags': ['graph:step:3'], 'metadata': {'langgraph_step': 3, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['tools'], 'langgraph_task_idx': 0}, 'data': {}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: should_continue\n",
      "{'event': 'on_chain_start', 'name': 'should_continue', 'run_id': '37869534-3f6f-4d14-a121-c18a4672612c', 'tags': ['seq:step:3'], 'metadata': {'langgraph_step': 3, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['tools'], 'langgraph_task_idx': 0}, 'data': {'input': {'messages': [HumanMessage(content=\"bing, What's the oldest parrot alive, and how much longer is that than the average lifespan of a parrot?\"), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_HXBAx6MHIpBrWoVaTZcCgWZw', 'function': {'arguments': '{\"query\":\"oldest parrot alive\"}', 'name': 'bing'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 18, 'prompt_tokens': 887, 'total_tokens': 905}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'tool_calls', 'logprobs': None, 'content_filter_results': {}}, id='run-a007cfce-535f-49b9-853d-d04a3e6abfe2-0', tool_calls=[{'name': 'bing', 'args': {'query': 'oldest parrot alive'}, 'id': 'call_HXBAx6MHIpBrWoVaTZcCgWZw'}], usage_metadata={'input_tokens': 887, 'output_tokens': 18, 'total_tokens': 905}), ToolMessage(content=\"The current oldest living parrot is **Poncho**, a Green-winged macaw, who is 92 years old. Poncho resides in the United States and was originally owned by Birds and Animals Unlimited. Most birds of Poncho's breed typically live to be around 50 or 60 years old, though some can reach 80 years with a healthy diet. Poncho has exceeded these expectations and is officially recognized as the world's oldest living parrot [[1]](https://www.oldest.org/animals/parrots/).\\n\\nIf you have any more questions or need further information, feel free to ask!\", name='bing', tool_call_id='call_HXBAx6MHIpBrWoVaTZcCgWZw'), AIMessage(content=\"The current oldest living parrot is **Poncho**, a Green-winged macaw, who is 92 years old. Poncho resides in the United States and was originally owned by Birds and Animals Unlimited. Most birds of Poncho's breed typically live to be around 50 or 60 years old, though some can reach 80 years with a healthy diet. Poncho has exceeded these expectations and is officially recognized as the world's oldest living parrot [[1]](https://www.oldest.org/animals/parrots/).\\n\\n### Comparison to Average Lifespan\\n- **Poncho's Age:** 92 years\\n- **Average Lifespan of Green-winged Macaw:** 50-60 years (with some reaching up to 80 years)\\n\\n**Difference:**\\n- Compared to the average lifespan of 50-60 years, Poncho has lived 32-42 years longer.\\n- Compared to the upper limit of 80 years, Poncho has lived 12 years longer.\", response_metadata={'token_usage': {'completion_tokens': 203, 'prompt_tokens': 1370, 'total_tokens': 1573}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'stop', 'logprobs': None, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}, id='run-f8e05920-fbbf-47d9-b313-5a9c140be0ef-0', usage_metadata={'input_tokens': 1370, 'output_tokens': 203, 'total_tokens': 1573})]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "The current oldest living parrot is **Poncho**, a Green-winged macaw, who is 92 years old. Poncho resides in the United States and was originally owned by Birds and Animals Unlimited. Most birds of Poncho's breed typically live to be around 50 or 60 years old, though some can reach 80 years with a healthy diet. Poncho has exceeded these expectations and is officially recognized as the world's oldest living parrot [[1]](https://www.oldest.org/animals/parrots/).\n",
       "\n",
       "### Comparison to Average Lifespan\n",
       "- **Poncho's Age:** 92 years\n",
       "- **Average Lifespan of Green-winged Macaw:** 50-60 years (with some reaching up to 80 years)\n",
       "\n",
       "**Difference:**\n",
       "- Compared to the average lifespan of 50-60 years, Poncho has lived 32-42 years longer.\n",
       "- Compared to the upper limit of 80 years, Poncho has lived 12 years longer."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "await print_events(\"bing, What's the oldest parrot alive, and how much longer is that than the average lifespan of a parrot?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b8df2e4-e808-47f6-b469-6ac586573904",
   "metadata": {},
   "source": [
    "### We can tell the app to return the results directly from the tool (without passing through the supervisor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fa1c8053-b1ef-4818-9076-13250472b327",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=======================\n",
      "\n",
      "Starting: LangGraph\n",
      "{'event': 'on_chain_start', 'run_id': 'db898ed0-50eb-4ca9-a65a-59bcd91d57ea', 'name': 'LangGraph', 'tags': [], 'metadata': {}, 'data': {'input': {'messages': [HumanMessage(content='docsearch, how diabetes affects covid? return the answer directly by setting return_direct=True')]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: supervisor\n",
      "{'event': 'on_chain_start', 'name': 'supervisor', 'run_id': 'aa4cec6a-1f62-40bf-b432-b45bb4ee80f0', 'tags': ['graph:step:1'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: should_continue\n",
      "{'event': 'on_chain_start', 'name': 'should_continue', 'run_id': '825dd388-34fe-491f-80f9-106251b02acc', 'tags': ['seq:step:3'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {'input': {'messages': [HumanMessage(content='docsearch, how diabetes affects covid? return the answer directly by setting return_direct=True'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_GRwEzTWeGQIlEl8fWDnph6Rd', 'function': {'arguments': '{\"query\":\"how diabetes affects covid?\",\"return_direct\":true}', 'name': 'docsearch'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 23, 'prompt_tokens': 880, 'total_tokens': 903}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'tool_calls', 'logprobs': None, 'content_filter_results': {}}, id='run-19c95aed-0610-4fbd-8f2a-38a7c8ae4225-0', tool_calls=[{'name': 'docsearch', 'args': {'query': 'how diabetes affects covid?', 'return_direct': True}, 'id': 'call_GRwEzTWeGQIlEl8fWDnph6Rd'}], usage_metadata={'input_tokens': 880, 'output_tokens': 23, 'total_tokens': 903})]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: tools_final\n",
      "{'event': 'on_chain_start', 'name': 'tools_final', 'run_id': 'ae1d803b-2ab5-487f-ab65-a5b2691bae6c', 'tags': ['graph:step:2'], 'metadata': {'langgraph_step': 2, 'langgraph_node': 'tools_final', 'langgraph_triggers': ['branch:supervisor:should_continue:tools_final'], 'langgraph_task_idx': 0}, 'data': {}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting tool: docsearch\n",
      "{'event': 'on_tool_start', 'name': 'docsearch', 'run_id': '9ef2be5f-8a4b-44cb-9bfc-c68bf8990557', 'tags': ['seq:step:1'], 'metadata': {'langgraph_step': 2, 'langgraph_node': 'tools_final', 'langgraph_triggers': ['branch:supervisor:should_continue:tools_final'], 'langgraph_task_idx': 0}, 'data': {'input': {'query': 'how diabetes affects covid?', 'return_direct': True}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "[{'supervisor': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_GRwEzTWeGQIlEl8fWDnph6Rd', 'function': {'arguments': '{\"query\":\"how diabetes affects covid?\",\"return_direct\":true}', 'name': 'docsearch'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 23, 'prompt_tokens': 880, 'total_tokens': 903}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'tool_calls', 'logprobs': None, 'content_filter_results': {}}, id='run-19c95aed-0610-4fbd-8f2a-38a7c8ae4225-0', tool_calls=[{'name': 'docsearch', 'args': {'query': 'how diabetes affects covid?', 'return_direct': True}, 'id': 'call_GRwEzTWeGQIlEl8fWDnph6Rd'}], usage_metadata={'input_tokens': 880, 'output_tokens': 23, 'total_tokens': 903})]}}, {'tools_final': {'messages': [ToolMessage(content='Diabetes significantly affects the progression and outcomes of COVID-19. Here are the key findings from various studies on this topic:\\n\\n1. **Increased Severity and Mortality Risk**:\\n   - Patients with diabetes are at a higher risk of severe COVID-19 and have a doubled mortality risk due to complications involving the pulmonary and cardiac systems [[2]](https://doi.org/10.1007/s00508-020-01672-3?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\\n\\n2. **Inflammatory Response and Complications**:\\n   - COVID-19 can exacerbate complications in individuals with diabetes through an imbalance in angiotensin-converting enzyme 2 (ACE2) activation pathways, leading to an inflammatory response. This imbalance can cause acute β-cell dysfunction in the pancreas, resulting in a hyperglycemic state [[1]](https://api.elsevier.com/content/article/pii/S1056872720303962; https://www.sciencedirect.com/science/article/pii/S1056872720303962?v=s5?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\\n\\n3. **Biomarkers and Hypercoagulable State**:\\n   - Diabetic patients with COVID-19 show higher levels of inflammatory biomarkers such as IL-6, C-reactive protein, serum ferritin, and D-dimer, indicating they are more susceptible to an inflammatory storm. This can lead to rapid deterioration of their condition [[3]](https://doi.org/10.1002/dmrr.3319; https://www.ncbi.nlm.nih.gov/pubmed/32233013/?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\\n\\n4. **Management Challenges**:\\n   - Managing diabetes during the pandemic is challenging due to limited access to outpatient clinics. This situation necessitates alternative treatment options like telemedicine. Proper glycemic control is crucial, and adjustments in antidiabetic drugs and insulin therapy may be required [[2]](https://doi.org/10.1007/s00508-020-01672-3?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\\n\\n5. **Preventive Measures**:\\n   - For individuals with diabetes, it is essential to stay hydrated, monitor blood glucose regularly, and check ketone bodies in urine if on insulin. Maintaining physical activity and a healthy diet are also crucial [[1]](https://api.elsevier.com/content/article/pii/S1056872720303962; https://www.sciencedirect.com/science/article/pii/S1056872720303962?v=s5?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\\n\\nIn summary, diabetes significantly increases the risk of severe COVID-19 and mortality due to various complications, including heightened inflammatory responses and hypercoagulable states. Effective management and preventive measures are crucial for mitigating these risks.', name='docsearch', tool_call_id='call_GRwEzTWeGQIlEl8fWDnph6Rd')]}}]\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "Diabetes significantly affects the progression and outcomes of COVID-19. Here are the key findings from various studies on this topic:\n",
       "\n",
       "1. **Increased Severity and Mortality Risk**:\n",
       "   - Patients with diabetes are at a higher risk of severe COVID-19 and have a doubled mortality risk due to complications involving the pulmonary and cardiac systems [[2]](https://doi.org/10.1007/s00508-020-01672-3?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\n",
       "\n",
       "2. **Inflammatory Response and Complications**:\n",
       "   - COVID-19 can exacerbate complications in individuals with diabetes through an imbalance in angiotensin-converting enzyme 2 (ACE2) activation pathways, leading to an inflammatory response. This imbalance can cause acute β-cell dysfunction in the pancreas, resulting in a hyperglycemic state [[1]](https://api.elsevier.com/content/article/pii/S1056872720303962; https://www.sciencedirect.com/science/article/pii/S1056872720303962?v=s5?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\n",
       "\n",
       "3. **Biomarkers and Hypercoagulable State**:\n",
       "   - Diabetic patients with COVID-19 show higher levels of inflammatory biomarkers such as IL-6, C-reactive protein, serum ferritin, and D-dimer, indicating they are more susceptible to an inflammatory storm. This can lead to rapid deterioration of their condition [[3]](https://doi.org/10.1002/dmrr.3319; https://www.ncbi.nlm.nih.gov/pubmed/32233013/?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\n",
       "\n",
       "4. **Management Challenges**:\n",
       "   - Managing diabetes during the pandemic is challenging due to limited access to outpatient clinics. This situation necessitates alternative treatment options like telemedicine. Proper glycemic control is crucial, and adjustments in antidiabetic drugs and insulin therapy may be required [[2]](https://doi.org/10.1007/s00508-020-01672-3?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\n",
       "\n",
       "5. **Preventive Measures**:\n",
       "   - For individuals with diabetes, it is essential to stay hydrated, monitor blood glucose regularly, and check ketone bodies in urine if on insulin. Maintaining physical activity and a healthy diet are also crucial [[1]](https://api.elsevier.com/content/article/pii/S1056872720303962; https://www.sciencedirect.com/science/article/pii/S1056872720303962?v=s5?sv=2022-11-02&ss=b&srt=sco&sp=rl&se=2026-07-09T13:52:04Z&st=2024-07-09T05:52:04Z&spr=https&sig=9Nx31tWOzf6CWylUZnGaciT9VDWVJSJ9vQulMcshm7Q%3D).\n",
       "\n",
       "In summary, diabetes significantly increases the risk of severe COVID-19 and mortality due to various complications, including heightened inflammatory responses and hypercoagulable states. Effective management and preventive measures are crucial for mitigating these risks."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "await print_events(\"docsearch, how diabetes affects covid? return the answer directly by setting return_direct=True\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ef806a5c-73e5-4cf4-848c-c52f8fe0f03f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=======================\n",
      "\n",
      "Starting: LangGraph\n",
      "{'event': 'on_chain_start', 'run_id': '881137da-5c50-4b72-868a-150d4d121bf4', 'name': 'LangGraph', 'tags': [], 'metadata': {}, 'data': {'input': {'messages': [HumanMessage(content='sqlsearch, How many people died of covid in Texas in 2020? return the answer directly by setting return_direct=True')]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: supervisor\n",
      "{'event': 'on_chain_start', 'name': 'supervisor', 'run_id': '1a18bdc6-f62c-48d4-a9a4-ac4b1d05f2c1', 'tags': ['graph:step:1'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: should_continue\n",
      "{'event': 'on_chain_start', 'name': 'should_continue', 'run_id': 'f7eacbbe-9d63-4217-b332-206b6f5c22bd', 'tags': ['seq:step:3'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {'input': {'messages': [HumanMessage(content='sqlsearch, How many people died of covid in Texas in 2020? return the answer directly by setting return_direct=True'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_X3ePb2cRJlq7ejIKg8uPoQFP', 'function': {'arguments': '{\"query\":\"How many people died of covid in Texas in 2020?\",\"return_direct\":true}', 'name': 'sqlsearch'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 31, 'prompt_tokens': 888, 'total_tokens': 919}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'tool_calls', 'logprobs': None, 'content_filter_results': {}}, id='run-02bbe1fb-0b1c-4a34-a4c5-58b0a1e05b28-0', tool_calls=[{'name': 'sqlsearch', 'args': {'query': 'How many people died of covid in Texas in 2020?', 'return_direct': True}, 'id': 'call_X3ePb2cRJlq7ejIKg8uPoQFP'}], usage_metadata={'input_tokens': 888, 'output_tokens': 31, 'total_tokens': 919})]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: tools_final\n",
      "{'event': 'on_chain_start', 'name': 'tools_final', 'run_id': '09c2a35e-ea2f-4ede-bd0f-2855f01819ca', 'tags': ['graph:step:2'], 'metadata': {'langgraph_step': 2, 'langgraph_node': 'tools_final', 'langgraph_triggers': ['branch:supervisor:should_continue:tools_final'], 'langgraph_task_idx': 0}, 'data': {}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting tool: sqlsearch\n",
      "{'event': 'on_tool_start', 'name': 'sqlsearch', 'run_id': '906794f2-606c-4aa5-bcf8-e34d7f17ff18', 'tags': ['seq:step:1'], 'metadata': {'langgraph_step': 2, 'langgraph_node': 'tools_final', 'langgraph_triggers': ['branch:supervisor:should_continue:tools_final'], 'langgraph_task_idx': 0}, 'data': {'input': {'query': 'How many people died of covid in Texas in 2020?', 'return_direct': True}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "[{'supervisor': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_X3ePb2cRJlq7ejIKg8uPoQFP', 'function': {'arguments': '{\"query\":\"How many people died of covid in Texas in 2020?\",\"return_direct\":true}', 'name': 'sqlsearch'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 31, 'prompt_tokens': 888, 'total_tokens': 919}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'tool_calls', 'logprobs': None, 'content_filter_results': {}}, id='run-02bbe1fb-0b1c-4a34-a4c5-58b0a1e05b28-0', tool_calls=[{'name': 'sqlsearch', 'args': {'query': 'How many people died of covid in Texas in 2020?', 'return_direct': True}, 'id': 'call_X3ePb2cRJlq7ejIKg8uPoQFP'}], usage_metadata={'input_tokens': 888, 'output_tokens': 31, 'total_tokens': 919})]}}, {'tools_final': {'messages': [ToolMessage(content=\"Final Answer: There were 2,841,253 people who died of covid in Texas in 2020.\\n\\nExplanation:\\nI queried the `covidtracking` table for the sum of the `death` column where the state is 'TX' and the date starts with '2020'. The query returned a total of 2,841,253 deaths for the year 2020. \\nI used the following query:\\n\\n```sql\\nSELECT SUM(death) AS total_deaths FROM covidtracking WHERE state = 'TX' AND date LIKE '2020%'\\n```\", name='sqlsearch', tool_call_id='call_X3ePb2cRJlq7ejIKg8uPoQFP')]}}]\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "Final Answer: There were 2,841,253 people who died of covid in Texas in 2020.\n",
       "\n",
       "Explanation:\n",
       "I queried the `covidtracking` table for the sum of the `death` column where the state is 'TX' and the date starts with '2020'. The query returned a total of 2,841,253 deaths for the year 2020. \n",
       "I used the following query:\n",
       "\n",
       "```sql\n",
       "SELECT SUM(death) AS total_deaths FROM covidtracking WHERE state = 'TX' AND date LIKE '2020%'\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "await print_events(\"sqlsearch, How many people died of covid in Texas in 2020? return the answer directly by setting return_direct=True\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "4af5d84f-59b9-43e3-8e1e-9fc55611ca7d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=======================\n",
      "\n",
      "Starting: LangGraph\n",
      "{'event': 'on_chain_start', 'run_id': '1ac95127-f520-4e2c-aa48-c04ca5336bcd', 'name': 'LangGraph', 'tags': [], 'metadata': {}, 'data': {'input': {'messages': [HumanMessage(content='Thank you!! you are a great assistant. BTW, what is my name?')]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: supervisor\n",
      "{'event': 'on_chain_start', 'name': 'supervisor', 'run_id': '1810a90f-b996-487c-bf8d-ddf708074de6', 'tags': ['graph:step:1'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n",
      "Starting: should_continue\n",
      "{'event': 'on_chain_start', 'name': 'should_continue', 'run_id': '1a9f8884-895c-4235-82bb-e2ecbb49b082', 'tags': ['seq:step:3'], 'metadata': {'langgraph_step': 1, 'langgraph_node': 'supervisor', 'langgraph_triggers': ['start:supervisor'], 'langgraph_task_idx': 0}, 'data': {'input': {'messages': [HumanMessage(content='Thank you!! you are a great assistant. BTW, what is my name?'), AIMessage(content=\"You're welcome! I'm glad to be of help. However, I don't have the ability to remember or know your name unless you've told me in this conversation. If you'd like, you can tell me your name now!\", response_metadata={'token_usage': {'completion_tokens': 44, 'prompt_tokens': 879, 'total_tokens': 923}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_abc28019ad', 'prompt_filter_results': [{'prompt_index': 0, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'jailbreak': {'filtered': False, 'detected': False}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}], 'finish_reason': 'stop', 'logprobs': None, 'content_filter_results': {'hate': {'filtered': False, 'severity': 'safe'}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': {'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}, id='run-2a55da81-704b-4c6e-b661-7912e9b7a53d-0', usage_metadata={'input_tokens': 879, 'output_tokens': 44, 'total_tokens': 923})]}}, 'parent_ids': []}\n",
      "\n",
      "=======================\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "You're welcome! I'm glad to be of help. However, I don't have the ability to remember or know your name unless you've told me in this conversation. If you'd like, you can tell me your name now!"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "await print_events(\"Thank you!! you are a great assistant. BTW, what is my name?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22270c98-6df9-47d1-ab76-3372e6f33cfb",
   "metadata": {},
   "source": [
    "# Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70919e9b-58d8-4899-baf8-5b9e3acb2cde",
   "metadata": {},
   "source": [
    "In this notebook, we have successfully implemented a multi-agent architecture utilizing LangGraph, marking a significant advancement in our capability to engineer complex, scalable systems. This framework enables the seamless integration of multiple agents, each performing distinct tasks, thereby facilitating more robust and intricate architectures.\n",
    "\n",
    "**Note: Currently, this architecture does not include a memory component**, which is crucial for tasks that require historical data recall or context retention over time. However, the development of this memory component is underway and promises to enhance the functionality significantly. Once integrated, this memory capability will allow graphs to perform more complex tasks that require understanding and analyzing past interactions, leading to more intelligent and adaptive responses.\n",
    "\n",
    "Stay tuned for updates!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5914b664-f094-4d03-82f0-03026f37edb1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10 - SDK v2",
   "language": "python",
   "name": "python310-sdkv2"
  },
  "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}