10_4_5_Using_Sub_Agents.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Appendix 10.4.5: Using Haiku as a sub-agent\n",
    "\n",
    "In this recipe, we'll demonstrate how to analyze Apple's 2023 financial earnings reports using Claude 3 Haiku sub-agent models to extract relevant information from earnings release PDFs. We'll then use Claude 3 Sonnet to generate a response to our question and create a graph using matplotlib to accompany its response."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Set up the environment\n",
    "First, let's install the required libraries and set up the Anthropic API client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pip install -qUr requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Import the required libraries\n",
    "import boto3\n",
    "import fitz\n",
    "from PIL import Image\n",
    "import io\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "import requests\n",
    "import os\n",
    "\n",
    "session = boto3.Session()\n",
    "region = session.region_name\n",
    "\n",
    "modelId = 'anthropic.claude-3-sonnet-20240229-v1:0'\n",
    "#modelId = 'anthropic.claude-3-haiku-20240307-v1:0'\n",
    "\n",
    "print(f'Using modelId: {modelId}')\n",
    "print('Using region: ', region)\n",
    "\n",
    "bedrock_client = boto3.client(service_name = 'bedrock-runtime', region_name = region,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Gather our documents and ask a question\n",
    "For this example, we will be using all Apple's financial statements from the 2023 financial year and asking about the net sales over the year."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List of Apple's earnings release PDF URLs\n",
    "pdf_urls = [\n",
    "    \"https://www.apple.com/newsroom/pdfs/fy2023-q4/FY23_Q4_Consolidated_Financial_Statements.pdf\",\n",
    "    \"https://www.apple.com/newsroom/pdfs/fy2023-q3/FY23_Q3_Consolidated_Financial_Statements.pdf\",\n",
    "    \"https://www.apple.com/newsroom/pdfs/FY23_Q2_Consolidated_Financial_Statements.pdf\",\n",
    "    \"https://www.apple.com/newsroom/pdfs/FY23_Q1_Consolidated_Financial_Statements.pdf\"\n",
    "]\n",
    "\n",
    "# User's question\n",
    "QUESTION = \"How did Apple's net sales change quarter to quarter in the 2023 financial year and what were the key contributors to the changes?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Download and convert PDFs to images\n",
    "Next, we'll define functions to download the earnings release PDFs and convert them to base64-encoded PNG images. We have to do this because these PDFs are full of tables that are hard to parse with traditional PDF parsers. It's easier if we just convert them to images and pass the images to Haiku.\n",
    "\n",
    "The ```download_pdf``` function downloads a PDF file from a given URL and saves it to the specified folder. The ```pdf_to_pngs``` function converts a PDF to a list of PNG images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Function to download a PDF file from a URL and save it to a specified folder\n",
    "def download_pdf(url, folder):\n",
    "    response = requests.get(url)\n",
    "    if response.status_code == 200:\n",
    "        file_name = os.path.join(folder, url.split(\"/\")[-1])\n",
    "        with open(file_name, \"wb\") as file:\n",
    "            file.write(response.content)\n",
    "        return file_name\n",
    "    else:\n",
    "        print(f\"Failed to download PDF from {url}\")\n",
    "        return None\n",
    "    \n",
    "# Define the function to convert a PDF to a list of base64-encoded PNG images\n",
    "def pdf_to_png(pdf_path, quality=75, max_size=(1024, 1024)):\n",
    "    # Open the PDF file\n",
    "    doc = fitz.open(pdf_path)\n",
    "    pdf_to_png_images = []\n",
    "\n",
    "    # Iterate through each page of the PDF\n",
    "    for page_num in range(doc.page_count):\n",
    "        # Load the page\n",
    "        page = doc.load_page(page_num)\n",
    "\n",
    "        # Render the page as a PNG image\n",
    "        pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))\n",
    "\n",
    "        # Convert the pixmap to a PIL Image\n",
    "        image = Image.frombytes(\"RGB\", [pix.width, pix.height], pix.samples)\n",
    "\n",
    "        # Resize the image if it exceeds the maximum size\n",
    "        if image.size[0] > max_size[0] or image.size[1] > max_size[1]:\n",
    "            image.thumbnail(max_size, Image.Resampling.LANCZOS)\n",
    "\n",
    "        # Convert the PIL image to bytes\n",
    "        image_data = io.BytesIO()\n",
    "        image.save(image_data, format='PNG', optimize=True, quality=quality)\n",
    "        image_data.seek(0)\n",
    "        pdf_to_png_image = image_data.getvalue()\n",
    "\n",
    "        # Append the PNG image bytes to the list\n",
    "        pdf_to_png_images.append(pdf_to_png_image)\n",
    "\n",
    "    # Close the PDF document\n",
    "    doc.close()\n",
    "\n",
    "    return pdf_to_png_images\n",
    "\n",
    "# Folder to save the downloaded PDFs\n",
    "folder = \"./images/using_sub_agents\"\n",
    "\n",
    "# Download the PDFs concurrently\n",
    "with ThreadPoolExecutor() as executor:\n",
    "    pdf_paths = list(executor.map(download_pdf, pdf_urls, [folder] * len(pdf_urls)))\n",
    "\n",
    "# Remove any None values (failed downloads) from pdf_paths\n",
    "pdf_paths = [path for path in pdf_paths if path is not None]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use ThreadPoolExecutor to download the PDFs concurrently and store the file paths in pdf_paths."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Generate a specific prompt for Haiku using Sonnet\n",
    "Let's use Opus as an orchestrator and have it write a specific prompt for each Haiku sub-agent based on the user provided question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_haiku_prompt(question):\n",
    "    prompt = f\"\"\"Based on the following question, please generate a specific prompt for an LLM sub-agent to extract relevant information from an earning's report PDF. Each sub-agent only has access to a single quarter's earnings report. Output only the prompt and nothing else.\\n\\nQuestion: {question}\"\"\"\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": 'user',\n",
    "            \"content\": [\n",
    "                {\"text\": prompt }\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "\n",
    "    converse_api_params = {\n",
    "        \"modelId\": modelId,\n",
    "        \"messages\": messages,\n",
    "    }\n",
    "\n",
    "    response = bedrock_client.converse(**converse_api_params)\n",
    "\n",
    "    return response['output']['message']['content'][0]['text']\n",
    "\n",
    "haiku_prompt = generate_haiku_prompt(QUESTION)\n",
    "print(haiku_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Extract information from PDFs\n",
    "Now, let's define our question and extract information from the PDFs using sub-agent Haiku models. We format the information from each model into a neatly defined set of XML tags."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def extract_info(pdf_path, haiku_prompt):\n",
    "    pdf_pngs = pdf_to_png(pdf_path)\n",
    "\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [\n",
    "                *[{\"image\": {\"format\": 'png', \"source\": {\"bytes\": pdf_png}}} for pdf_png in pdf_pngs],\n",
    "                {\"text\": haiku_prompt}\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "\n",
    "    converse_api_params = {\n",
    "        \"modelId\": \"anthropic.claude-3-haiku-20240307-v1:0\",\n",
    "        \"messages\": messages,\n",
    "    }\n",
    "    response = bedrock_client.converse(**converse_api_params)\n",
    "\n",
    "    return response['output']['message']['content'][0]['text'], pdf_path\n",
    "\n",
    "def process_pdf(pdf_path):\n",
    "    return extract_info(pdf_path, QUESTION)\n",
    "\n",
    "# Process the PDFs concurrently with Haiku sub-agent models\n",
    "with ThreadPoolExecutor() as executor:\n",
    "    extracted_info_list = list(executor.map(process_pdf, pdf_paths))\n",
    "\n",
    "extracted_info = \"\"\n",
    "# Display the extracted information from each model call\n",
    "for info in extracted_info_list:\n",
    "    extracted_info += \"<info quarter=\\\"\" + info[1].split(\"/\")[-1].split(\"_\")[1] + \"\\\">\" + info[0] + \"</info>\\n\"\n",
    "print(extracted_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We extract information from the PDFs concurrently using sub-agent models and combine the extracted information. We then prepare the messages for the powerful model, including the question and the extracted information, and ask it to generate a response and matplotlib code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Pass the information to Sonnet to generate a response\n",
    "Now that we have fetched the information from each PDF using the sub-agents, let's call Opus to actually answer the question and write code to create a graph to accompany the answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Prepare the messages for the powerful model\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\"text\": f\"Based on the following extracted information from Apple's earnings releases, please provide a response to the question: {QUESTION}\\n\\nAlso, please generate Python code using the matplotlib library to accompany your response. Enclose the code within <code> tags.\\n\\nExtracted Information:\\n{extracted_info}\"}\n",
    "        ]\n",
    "    }\n",
    "]\n",
    "\n",
    "# Generate the matplotlib code using the powerful model\n",
    "converse_api_params = {\n",
    "    \"modelId\": \"anthropic.claude-3-sonnet-20240229-v1:0\",\n",
    "    \"messages\": messages,\n",
    "    \"inferenceConfig\": {\"maxTokens\": 4096},\n",
    "}\n",
    "response = bedrock_client.converse(**converse_api_params)\n",
    "\n",
    "generated_response = response['output']['message']['content'][0]['text']\n",
    "print(\"Generated Response:\")\n",
    "print(generated_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Extract response and execute Matplotlib code\n",
    "Finally, let's extract the matplotlib code from the generated response and execute it to visualize the revenue growth trend.\n",
    "\n",
    "We define the ```extract_code_and_response``` function to extract the matplotlib code and non-code response from the generated response. We print the non-code response and execute the matplotlib code if it is found.\n",
    "\n",
    "Note that it is not good practice to use ```exec``` on model-written code outside of a sandbox but for the purposes of this demo we are doing it :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Extract the matplotlib code from the response\n",
    "# Function to extract the code and non-code parts from the response\n",
    "def extract_code_and_response(response):\n",
    "    start_tag = \"<code>\"\n",
    "    end_tag = \"</code>\"\n",
    "    start_index = response.find(start_tag)\n",
    "    end_index = response.find(end_tag)\n",
    "    if start_index != -1 and end_index != -1:\n",
    "        code = response[start_index + len(start_tag):end_index].strip()\n",
    "        non_code_response = response[:start_index].strip()\n",
    "        return code, non_code_response\n",
    "    else:\n",
    "        return None, response.strip()\n",
    "\n",
    "matplotlib_code, non_code_response = extract_code_and_response(generated_response)\n",
    "\n",
    "print(non_code_response)\n",
    "if matplotlib_code:\n",
    "\n",
    "    # Execute the extracted matplotlib code\n",
    "    exec(matplotlib_code)\n",
    "else:\n",
    "    print(\"No matplotlib code found in the response.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}