3 - rewoo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup Variables\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "aws_profile_name = \"dev\"\n",
    "aws_region = \"us-east-1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.globals import set_debug, set_verbose\n",
    "\n",
    "set_debug(False)\n",
    "set_verbose(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup AWS Bedrock Model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "from langchain_aws import ChatBedrockConverse\n",
    "\n",
    "boto3.setup_default_session(profile_name=aws_profile_name)\n",
    "BEDROCK_CLIENT = boto3.client(\"bedrock-runtime\", aws_region)\n",
    "model = ChatBedrockConverse(\n",
    "    client=BEDROCK_CLIENT,\n",
    "    model=\"anthropic.claude-3-5-sonnet-20240620-v1:0\",\n",
    "    max_tokens=1000,\n",
    "    temperature=0.2,\n",
    "    top_p=0.9,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Our Tools\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get Teammates By Name\n",
    "\n",
    "Go to our \"api\" to get all teammates that have similar names. In this case, our API is actually just a JSON file we read in at runtime.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.tools import tool\n",
    "from pydantic import BaseModel, TypeAdapter\n",
    "\n",
    "\n",
    "class Teammate(BaseModel):\n",
    "    id: int\n",
    "    name: str\n",
    "    favorite_color: str\n",
    "\n",
    "\n",
    "@tool(parse_docstring=True)\n",
    "def get_teammates_by_name(search_str: str) -> list[Teammate]:\n",
    "    \"\"\"Call to an API to get teammate details searching by their name.\n",
    "\n",
    "    Args:\n",
    "        search_str: The search string used to look for a teammate.\n",
    "            This can be a partial string and will return all teammates\n",
    "            who contain this search string as a substring, case insensitive.\n",
    "\n",
    "    Returns:\n",
    "        list[Teammate]: A list of any teammates with names containing the search string.\n",
    "    \"\"\"\n",
    "    with open(\"data/teammates.json\", \"r\") as file:\n",
    "        teammates = TypeAdapter(list[Teammate]).validate_json(file.read())\n",
    "    return [\n",
    "        teammate\n",
    "        for teammate in teammates\n",
    "        if search_str.lower() in teammate.name.lower()\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Consolidate All Our Tools\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import cast\n",
    "\n",
    "from langchain_core.tools.simple import Tool\n",
    "\n",
    "tools = cast(list[Tool], [get_teammates_by_name])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# REWOO Parts\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Plan\n",
    "\n",
    "First thing we need to do is generate a plan to solve our problem. This requires more work than our continuous loop technique, but results in fewer calls and a more definite set of actions you can communicate to the user.\n",
    "\n",
    "Lets define what we want a \"Plan\" object to look like. We are going to leverage Pydantic descriptions so we don't have to define this information in our prompt.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any\n",
    "\n",
    "from langchain.output_parsers import OutputFixingParser\n",
    "from langchain_core.output_parsers.pydantic import PydanticOutputParser\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "class Step(BaseModel):\n",
    "    id: int = Field(\n",
    "        description=\"Unique identifier for the step. IDs must be unique and in order.\"\n",
    "    )\n",
    "    summary: str = Field(\n",
    "        description=\"Summary of the step in plain english for a human to read.\"\n",
    "    )\n",
    "    tool_name: str = Field(description=\"Name of the tool to be used in this step.\")\n",
    "    tool_input: str = Field(description=\"Input to the tool.\")\n",
    "    tool_output: Any = Field(\n",
    "        None,\n",
    "        description=\"Output from the tool, can be null, a string, or an object with unknown properties. Default is null.\",\n",
    "    )\n",
    "\n",
    "\n",
    "class Plan(BaseModel):\n",
    "    summary: str = Field(\n",
    "        description=\"A summary of your plan in plain english for a human to read.\"\n",
    "    )\n",
    "    steps: list[Step] = Field(\n",
    "        description=\"An ordered list of Step objects to complete the task.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Also, lets quickly create some parsers that we can use later down the line for both prompting and failover\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The output should be formatted as a JSON instance that conforms to the JSON schema below.\n",
      "\n",
      "As an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]}\n",
      "the object {\"foo\": [\"bar\", \"baz\"]} is a well-formatted instance of the schema. The object {\"properties\": {\"foo\": [\"bar\", \"baz\"]}} is not well-formatted.\n",
      "\n",
      "Here is the output schema:\n",
      "```\n",
      "{\"$defs\": {\"Step\": {\"properties\": {\"id\": {\"description\": \"Unique identifier for the step. IDs must be unique and in order.\", \"title\": \"Id\", \"type\": \"integer\"}, \"summary\": {\"description\": \"Summary of the step in plain english for a human to read.\", \"title\": \"Summary\", \"type\": \"string\"}, \"tool_name\": {\"description\": \"Name of the tool to be used in this step.\", \"title\": \"Tool Name\", \"type\": \"string\"}, \"tool_input\": {\"description\": \"Input to the tool.\", \"title\": \"Tool Input\", \"type\": \"string\"}, \"tool_output\": {\"default\": null, \"description\": \"Output from the tool, can be null, a string, or an object with unknown properties. Default is null.\", \"title\": \"Tool Output\"}}, \"required\": [\"id\", \"summary\", \"tool_name\", \"tool_input\"], \"title\": \"Step\", \"type\": \"object\"}}, \"properties\": {\"summary\": {\"description\": \"A summary of your plan in plain english for a human to read.\", \"title\": \"Summary\", \"type\": \"string\"}, \"steps\": {\"description\": \"An ordered list of Step objects to complete the task.\", \"items\": {\"$ref\": \"#/$defs/Step\"}, \"title\": \"Steps\", \"type\": \"array\"}}, \"required\": [\"summary\", \"steps\"]}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "plan_parser = PydanticOutputParser(pydantic_object=Plan)\n",
    "print(plan_parser.get_format_instructions())\n",
    "parser = OutputFixingParser.from_llm(parser=plan_parser, llm=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Our Prompt to Generate Our Plan\n",
    "\n",
    "We get to leverage Langchain parser and tools so we only have to manage definitions in once place.\n",
    "\n",
    "Examples are good to throw in so it finds patterns to latch on to.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "planning_prompt = f\"\"\"\n",
    "A human is going to provide you with a task or message, make plans that can solve their problem step by step.\n",
    "\n",
    "{plan_parser.get_format_instructions().replace(\"{\",\"{{\").replace(\"}\",\"}}\")},\n",
    "\n",
    "When creating the plans, if you need something from a previous step you can reference it directly using '[stepId]'.\n",
    "For example, if you need the output from step 1, you can reference it as '[1]'. If you need the output from step 2, you can reference it as '[2]', and so on.\n",
    "\n",
    "For each plan, indicate which external tool together with tool input to retrieve evidence. You do not have to use any tools if you do not need to.\n",
    "Make sure the plan is optimized and as few of steps as required to perform the given task.\n",
    "You also have access to the history of the conversation with the user to help you plan effectively.\n",
    "\n",
    "Here are some examples of the JSON schema with the user input and the task to be completed including available tools:\n",
    "<examples>\n",
    "Input: Are there multiple teammates with the name Lovelace?\n",
    "Output Plan:\n",
    "{{{{\n",
    "  'summary': 'We need to get all teammates with the name Lovelace',\n",
    "  'steps': [\n",
    "    {{{{\n",
    "      'id': 1,\n",
    "      'summary': 'Search the teammates API for all teammates with the name Lovelace.',\n",
    "      'tool_name': 'get_teammates_by_name',\n",
    "      'tool_input': 'lovelace',\n",
    "      'tool_output': null\n",
    "    }}}}\n",
    "  ]\n",
    "}}}}\n",
    "\n",
    "Input: 5x + 3 = 18, solve for x.\n",
    "Output Plan:\n",
    "{{{{\n",
    "  'summary': 'Solve the equation 5x + 3 = 18 for x, we do not need any tools for this task.',\n",
    "  'steps': []\n",
    "}}}}\n",
    "\n",
    "Input: Hello\n",
    "Output Plan:\n",
    "{{{{\n",
    "  'summary': 'The user is just saying hello, we do not need any tools for this task.',\n",
    "  'steps': []\n",
    "}}}}\n",
    "</examples>\n",
    "\n",
    "Make sure to follow the schema strictly and provide the plan in the correct format. \n",
    "You will only respond with valid JSON.\n",
    "Do not include any additional information or text in the response as it will be considered as an invalid response and you will be asked to provide the answer again.\n",
    "You will never change the user's task or provide an answer to the task.\n",
    "You will only provide a plan to solve the task.\n",
    "DO NOT call tools that are not provided to you.\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example of a getting a plan\n",
    "\n",
    "NOTE: Sometimes we MAY get calls out to tools that do not exist. We should try and handle those gracefully. Below, I simply ignore the calls.\n",
    "\n",
    "This could potentially be avoided by providing better tool info in our prompt, but for now, we are going to rely on what langchain is doing naturally during the `bind_tools` call.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"summary\": \"To determine if Matt and Grey have the same favorite color, we need to search for information about both teammates.\",\n",
      "  \"steps\": [\n",
      "    {\n",
      "      \"id\": 1,\n",
      "      \"summary\": \"Search for teammates with the name Matt.\",\n",
      "      \"tool_name\": \"get_teammates_by_name\",\n",
      "      \"tool_input\": \"Matt\",\n",
      "      \"tool_output\": null\n",
      "    },\n",
      "    {\n",
      "      \"id\": 2,\n",
      "      \"summary\": \"Search for teammates with the name Grey.\",\n",
      "      \"tool_name\": \"get_teammates_by_name\",\n",
      "      \"tool_input\": \"Grey\",\n",
      "      \"tool_output\": null\n",
      "    }\n",
      "  ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", planning_prompt),\n",
    "        (\"human\", \"Do Matt and Grey have the same favorite color?\"),\n",
    "    ]\n",
    ")\n",
    "planning_chain = prompt | model.bind_tools(tools) | parser\n",
    "plan_output: Plan = planning_chain.invoke({})\n",
    "print(plan_output.model_dump_json(indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving\n",
    "\n",
    "After our `Planning` and `Tool Calling` phases, we need to wrap everything up with a `Solve` phase. We will feed in the plan with all the outputs attached.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "solve_prompt = \"\"\"\n",
    "Solve the following task or problem.\n",
    "To solve the problem, we have made step-by-step Plan and retrieved corresponding output for each step.\n",
    "Use them with caution since long output might contain irrelevant information.\n",
    "\n",
    "<plan>\n",
    "{plan}\n",
    "</plan>\n",
    "\n",
    "Now solve the question or task according to provided plan and output above.\n",
    "Respond with the answer directly with no extra words unless they provide helpful detail.\n",
    "If you don't know the answer, don't make something up,just tell what information you need to solve the problem or answer the message.\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Assembling our Graph\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Overall we will want our final graph to do these things:\n",
    "\n",
    "1. Make a plan\n",
    "2. Iterate over steps in that plan and call tools until all tools have an `output`\n",
    "3. Solve\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First, State Definition and Some Helper Methods\n",
    "\n",
    "This defines what state we will be passing around our graph every time we need to execute or make a decision of where to go next.\n",
    "\n",
    "In this case, we simply want to use the built in `MessagesState`, but add our plan object to it as well.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import MessagesState\n",
    "\n",
    "\n",
    "class RewooState(MessagesState):\n",
    "    plan: Plan\n",
    "\n",
    "\n",
    "def get_next_step(state: RewooState):\n",
    "    return next(\n",
    "        (step for step in state[\"plan\"].steps if step.tool_output is None), None\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now the Execution Methods. These will make our graph `nodes`.\n",
    "\n",
    "All of these methods will receive our graph state and return and object that will be merged with our graph state. This can just be our modified full state again if that is easier.\n",
    "\n",
    "NOTE: There is some code in our `call_tool` method that tries to reconcile prior tools outputs that are referenced in future tool inputs. This could use some love.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "\n",
    "def create_plan(state: RewooState):\n",
    "    print(\"Planning\")\n",
    "    print(state)\n",
    "    prompt = ChatPromptTemplate.from_messages(\n",
    "        [(\"system\", planning_prompt), *state[\"messages\"]]\n",
    "    )\n",
    "    planning_chain = prompt | model.bind_tools(tools) | parser\n",
    "    plan: Plan = planning_chain.invoke({})\n",
    "    return {\"plan\": plan}\n",
    "\n",
    "\n",
    "def call_tool(state: RewooState):\n",
    "    print(\"Calling Tool\")\n",
    "    print(state)\n",
    "    step = get_next_step(state)\n",
    "    if step:\n",
    "        tool = next((tool for tool in tools if tool.name == step.tool_name), None)\n",
    "        if tool:\n",
    "\n",
    "            def get_tool_output(match_obj):\n",
    "                return str(\n",
    "                    next(\n",
    "                        (\n",
    "                            step\n",
    "                            for step in state[\"plan\"].steps\n",
    "                            if step.id == match_obj.group(1)\n",
    "                        )\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            step.tool_input = re.sub(r\"\\[(\\d+)\\]\", get_tool_output, step.tool_input)\n",
    "            step.tool_output = tool.invoke(step.tool_input)\n",
    "        else:\n",
    "            step.tool_output = \"Tool does not exist.\"\n",
    "    return state\n",
    "\n",
    "\n",
    "def solve(state: RewooState):\n",
    "    print(\"Solving\")\n",
    "    print(state)\n",
    "    prompt = ChatPromptTemplate.from_messages(\n",
    "        [(\"system\", solve_prompt), *state[\"messages\"]]\n",
    "    )\n",
    "    solve_chain = prompt | model\n",
    "    result = solve_chain.invoke(state)\n",
    "    result.additional_kwargs[\"plan\"] = state[\"plan\"]\n",
    "    return {\"messages\": [result]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now Decision Making/Routing Methods. These will make our graph `edges`.\n",
    "\n",
    "All of these methods will receive our graph state object and return a string that represents the next node that should be executed.\n",
    "\n",
    "Routing methods are only needed when there is a `conditional` routing event. Some events you will see later on are simply hard coded since they always follow each other.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def should_call_tool_or_solve(state: RewooState):\n",
    "    return call_tool.__name__ if get_next_step(state) else solve.__name__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now we put it all together!\n",
    "\n",
    "NOTE: The ascii graph is not correct due to ENUM vs Str and not great Python typing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "        +-----------+     \n",
      "        | __start__ |     \n",
      "        +-----------+     \n",
      "              *           \n",
      "              *           \n",
      "              *           \n",
      "       +-------------+    \n",
      "       | create_plan |    \n",
      "       +-------------+    \n",
      "              .           \n",
      "              .           \n",
      "              .           \n",
      "        +-----------+     \n",
      "        | call_tool |     \n",
      "        +-----------+     \n",
      "         ..        ..     \n",
      "       ..            .    \n",
      "      .               ..  \n",
      "+-------+               . \n",
      "| solve |             ..  \n",
      "+-------+            .    \n",
      "         **        ..     \n",
      "           **    ..       \n",
      "             *  .         \n",
      "         +---------+      \n",
      "         | __end__ |      \n",
      "         +---------+      \n"
     ]
    }
   ],
   "source": [
    "from langgraph.graph import StateGraph\n",
    "\n",
    "graph = StateGraph(RewooState)\n",
    "\n",
    "# definitions\n",
    "graph.add_node(create_plan)\n",
    "graph.add_node(call_tool)\n",
    "graph.add_node(solve)\n",
    "\n",
    "# navigation\n",
    "graph.set_entry_point(create_plan.__name__)\n",
    "graph.add_edge(create_plan.__name__, call_tool.__name__)\n",
    "graph.add_conditional_edges(call_tool.__name__, should_call_tool_or_solve)\n",
    "graph.set_finish_point(solve.__name__)\n",
    "\n",
    "agent = graph.compile()\n",
    "agent.get_graph().print_ascii()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now that we have our agent assembled, we just need to seed it with some messages and get back results!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Planning\n",
      "{'messages': [HumanMessage(content='Do Matt and Grey have the same favorite color?', additional_kwargs={}, response_metadata={}, id='259ddf37-c3c4-4efa-8987-5705e8967faf')]}\n",
      "Calling Tool\n",
      "{'messages': [HumanMessage(content='Do Matt and Grey have the same favorite color?', additional_kwargs={}, response_metadata={}, id='259ddf37-c3c4-4efa-8987-5705e8967faf')], 'plan': Plan(summary='To determine if Matt and Grey have the same favorite color, we need to search for their information in the teammates database and compare their favorite colors.', steps=[Step(id=1, summary=\"Search for Matt's information in the teammates database.\", tool_name='get_teammates_by_name', tool_input='Matt', tool_output=None), Step(id=2, summary=\"Search for Grey's information in the teammates database.\", tool_name='get_teammates_by_name', tool_input='Grey', tool_output=None)])}\n",
      "Calling Tool\n",
      "{'messages': [HumanMessage(content='Do Matt and Grey have the same favorite color?', additional_kwargs={}, response_metadata={}, id='259ddf37-c3c4-4efa-8987-5705e8967faf')], 'plan': Plan(summary='To determine if Matt and Grey have the same favorite color, we need to search for their information in the teammates database and compare their favorite colors.', steps=[Step(id=1, summary=\"Search for Matt's information in the teammates database.\", tool_name='get_teammates_by_name', tool_input='Matt', tool_output=[Teammate(id=4, name='Matt Vincent', favorite_color='green')]), Step(id=2, summary=\"Search for Grey's information in the teammates database.\", tool_name='get_teammates_by_name', tool_input='Grey', tool_output=None)])}\n",
      "Solving\n",
      "{'messages': [HumanMessage(content='Do Matt and Grey have the same favorite color?', additional_kwargs={}, response_metadata={}, id='259ddf37-c3c4-4efa-8987-5705e8967faf')], 'plan': Plan(summary='To determine if Matt and Grey have the same favorite color, we need to search for their information in the teammates database and compare their favorite colors.', steps=[Step(id=1, summary=\"Search for Matt's information in the teammates database.\", tool_name='get_teammates_by_name', tool_input='Matt', tool_output=[Teammate(id=4, name='Matt Vincent', favorite_color='green')]), Step(id=2, summary=\"Search for Grey's information in the teammates database.\", tool_name='get_teammates_by_name', tool_input='Grey', tool_output=[Teammate(id=1, name='Grey Lovelace', favorite_color='purple')])])}\n"
     ]
    }
   ],
   "source": [
    "output = agent.invoke(\n",
    "    {\"messages\": [(\"human\", \"Do Matt and Grey have the same favorite color?\")]}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Final Answer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No, Matt and Grey do not have the same favorite color. Matt's favorite color is green, while Grey's favorite color is purple.\n"
     ]
    }
   ],
   "source": [
    "print(output[\"messages\"][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Message Details\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"content\": \"Do Matt and Grey have the same favorite color?\",\n",
      "  \"additional_kwargs\": {},\n",
      "  \"response_metadata\": {},\n",
      "  \"type\": \"human\",\n",
      "  \"name\": null,\n",
      "  \"id\": \"259ddf37-c3c4-4efa-8987-5705e8967faf\",\n",
      "  \"example\": false\n",
      "}\n",
      "{\n",
      "  \"content\": \"No, Matt and Grey do not have the same favorite color. Matt's favorite color is green, while Grey's favorite color is purple.\",\n",
      "  \"additional_kwargs\": {\n",
      "    \"plan\": {\n",
      "      \"summary\": \"To determine if Matt and Grey have the same favorite color, we need to search for their information in the teammates database and compare their favorite colors.\",\n",
      "      \"steps\": [\n",
      "        {\n",
      "          \"id\": 1,\n",
      "          \"summary\": \"Search for Matt's information in the teammates database.\",\n",
      "          \"tool_name\": \"get_teammates_by_name\",\n",
      "          \"tool_input\": \"Matt\",\n",
      "          \"tool_output\": [\n",
      "            {\n",
      "              \"id\": 4,\n",
      "              \"name\": \"Matt Vincent\",\n",
      "              \"favorite_color\": \"green\"\n",
      "            }\n",
      "          ]\n",
      "        },\n",
      "        {\n",
      "          \"id\": 2,\n",
      "          \"summary\": \"Search for Grey's information in the teammates database.\",\n",
      "          \"tool_name\": \"get_teammates_by_name\",\n",
      "          \"tool_input\": \"Grey\",\n",
      "          \"tool_output\": [\n",
      "            {\n",
      "              \"id\": 1,\n",
      "              \"name\": \"Grey Lovelace\",\n",
      "              \"favorite_color\": \"purple\"\n",
      "            }\n",
      "          ]\n",
      "        }\n",
      "      ]\n",
      "    }\n",
      "  },\n",
      "  \"response_metadata\": {\n",
      "    \"ResponseMetadata\": {\n",
      "      \"RequestId\": \"8290c376-1b69-4cb3-a442-cd13b7350aac\",\n",
      "      \"HTTPStatusCode\": 200,\n",
      "      \"HTTPHeaders\": {\n",
      "        \"date\": \"Thu, 24 Oct 2024 22:14:01 GMT\",\n",
      "        \"content-type\": \"application/json\",\n",
      "        \"content-length\": \"309\",\n",
      "        \"connection\": \"keep-alive\",\n",
      "        \"x-amzn-requestid\": \"8290c376-1b69-4cb3-a442-cd13b7350aac\"\n",
      "      },\n",
      "      \"RetryAttempts\": 0\n",
      "    },\n",
      "    \"stopReason\": \"end_turn\",\n",
      "    \"metrics\": {\n",
      "      \"latencyMs\": 1032\n",
      "    }\n",
      "  },\n",
      "  \"type\": \"ai\",\n",
      "  \"name\": null,\n",
      "  \"id\": \"run-a06727eb-32ca-459e-9637-0b3198827601-0\",\n",
      "  \"example\": false,\n",
      "  \"tool_calls\": [],\n",
      "  \"invalid_tool_calls\": [],\n",
      "  \"usage_metadata\": {\n",
      "    \"input_tokens\": 289,\n",
      "    \"output_tokens\": 31,\n",
      "    \"total_tokens\": 320\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "for message in output[\"messages\"]:\n",
    "    print(json.dumps(message, indent=2, default=lambda x: x.dict()))"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}