{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Build Your Own AI Podcasting Agent with LangGraph, Gemini, and Chirp 3\n",
        "## AI-Powered Podcast Creation with Automated Research, Writing, and Refinement\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/orchestration/langgraph_gemini_podcast.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Forchestration%2Flanggraph_gemini_podcast.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/orchestration/langgraph_gemini_podcast.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/orchestration/langgraph_gemini_podcast.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/orchestration/langgraph_gemini_podcast.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/orchestration/langgraph_gemini_podcast.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/orchestration/langgraph_gemini_podcast.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/orchestration/langgraph_gemini_podcast.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/orchestration/langgraph_gemini_podcast.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| Author |\n",
        "| --- |\n",
        "| [Kristopher Overholt](https://github.com/koverholt/) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Creating a podcast can be a very involved process, requiring extensive research, writing, editing, and production. **What if there was a way to leverage the power of AI to streamline the creation of a podcast, automating many of the tasks traditionally performed by humans?** [NotebookLM](https://notebooklm.google.com/), for example, lets users easily generate [audio overviews based on documents](https://blog.google/technology/ai/notebooklm-audio-overviews/).\n",
        "\n",
        "#### 🔈🔈 [Listen to a sample podcast generated by this notebook!](https://storage.googleapis.com/github-repo/generative-ai/gemini/orchestration/langgraph/gemini-podcast.mp3) 🔈🔈\n",
        "\n",
        "But what if you want to customize the length of the podcast, the voices, or the conversation flow and augment it with additional research tools? In this notebook, **you'll recreate this kind of podcast generation functionality by building an AI agent to do the heavy lifting and then customize the entire flow yourself!**\n",
        "\n",
        "**This notebook demonstrates how to build a [LangGraph](https://langchain-ai.github.io/langgraph/)-powered AI agent to research, write, and refine a podcast script using the [Gemini API in in Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models).** You'll use LangGraph and LangChain to orchestrate calls to Gemini along with calls to different search tools, allowing the AI to learn about a given topic before writing about it. Then, the AI will critique its work and iterate on the podcast script, improving it with each revision.\n",
        "\n",
        "Here's how you'll build and use our AI podcasting agent:\n",
        "\n",
        "- **[User]** Define the podcast topic: Provide a clear and concise topic for the podcast.\n",
        "- **[Agent]** Generate an outline: Use Gemini to create a high-level outline, structuring the podcast's flow.\n",
        "- **[Agent]** Conduct research: The AI agent will use search tools like arXiv, PubMed, and Wikipedia to gather relevant information.\n",
        "- **[Agent]** Write a script: Gemini will generate an engaging podcast script, incorporating the research findings.\n",
        "- **[Agent]** Critique and iterate: The agent will analyze its script, provide a critique, then generate a revised draft.\n",
        "- **[Agent]** Generate audio: You'll use text-to-speech to generate audio for each line of the podcast script.\n",
        "\n",
        "<img src=\"https://storage.googleapis.com/github-repo/generative-ai/gemini/orchestration/langgraph/gemini-podcast-agent.jpg\" width=\"400px\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started\n",
        "\n",
        "This section sets up the environment for the AI podcast agent. This includes:\n",
        "\n",
        "- **Installing Libraries:**  Installing the required Python libraries\n",
        "- **Restarting Runtime (Colab Only):**  Restarting the Colab runtime\n",
        "- **Authenticating Environment (Colab Only):**  Authenticating to Google Cloud\n",
        "- **Setting Project Information:**  Setting up your Google Cloud project"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68664edd8334"
      },
      "source": [
        "### Install FFMpeg on your machine\n",
        "\n",
        "Install the FFMpeg libraries for running the AudioSegment library.",
        "\n",
        "If you're using Vertex AI, open the **Terminal**. Run the following command:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "outputs": [],
      "source": [
        "sudo apt-get update && apt-get install ffmpeg -y"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages\n",
        "\n",
        "This code cell installs the necessary Python libraries for running the AI podcast agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install -q -U \\\n",
        "    arxiv \\\n",
        "    google-cloud-aiplatform \\\n",
        "    google-cloud-texttospeech \\\n",
        "    langgraph \\\n",
        "    langchain-google-vertexai \\\n",
        "    langchain-community \\\n",
        "    pydub \\\n",
        "    pymupdf \\\n",
        "    wikipedia \\\n",
        "    xmltodict \\\n",
        "    numpy==1.24.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import vertexai\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\", isTemplate: true}\n",
        "if PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Building the AI podcasting agent\n",
        "\n",
        "This section constructs the AI agent. Key steps include:\n",
        "\n",
        "- **Initializing Agent Memory and State:** Setting up the agent's memory and defining its data structure\n",
        "- **Initializing the Gemini Model:**  Loading the Gemini language model from Vertex AI\n",
        "- **Defining Search Tools:**  Creating tools to access information sources like arXiv, PubMed, and Wikipedia\n",
        "- **Defining Workflow Stages:** Defining each stage of the workflow, including prompts and functions\n",
        "- **Compiling the Workflow:**  Structuring the workflow as a graph using LangGraph"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "This section imports the necessary libraries for LangGraph, LangChain, Vertex AI, and other utilities needed for your agent's functionality.\n",
        "\n",
        "This includes tools for interacting with the Gemini API, defining custom tools, managing agent state, and displaying results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6fc324893334"
      },
      "outputs": [],
      "source": [
        "# Common libraries\n",
        "import logging\n",
        "import os\n",
        "import re\n",
        "\n",
        "# Typing utilities for data validation and schema definitions\n",
        "from typing import TypedDict\n",
        "\n",
        "from IPython.display import Audio, Image\n",
        "\n",
        "# Libraries for text-to-speech generation and audio processing\n",
        "from google.cloud import texttospeech\n",
        "from langchain.schema.document import Document\n",
        "\n",
        "# Tools\n",
        "from langchain_community.retrievers import (\n",
        "    ArxivRetriever,\n",
        "    PubMedRetriever,\n",
        "    WikipediaRetriever,\n",
        ")\n",
        "\n",
        "# LangChain and LangGraph components for message handling and tool integration\n",
        "from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage\n",
        "from langchain_core.tools import tool\n",
        "\n",
        "# LangChain integrations for Gemini API in Google AI Studio and Vertex AI\n",
        "from langchain_google_vertexai import ChatVertexAI\n",
        "from langgraph.checkpoint.memory import MemorySaver\n",
        "from langgraph.graph import END, StateGraph\n",
        "from langgraph.prebuilt import ToolNode\n",
        "from pydub import AudioSegment\n",
        "\n",
        "# Set logging level to ERROR to filter warnings\n",
        "logger = logging.getLogger()\n",
        "logger.setLevel(logging.ERROR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e43229f3ad4f"
      },
      "source": [
        "### Initialize agent memory and agent state\n",
        "\n",
        "Here, you initialize your [agent's memory](https://langchain-ai.github.io/langgraph/how-tos/memory/manage-conversation-history/) to store information during the workflow.\n",
        "\n",
        "You also define the schema for your [agent's state](https://langchain-ai.github.io/langgraph/how-tos/state-model/), which includes the podcast outline, search queries, and drafts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf93d5f0ce00"
      },
      "outputs": [],
      "source": [
        "# Initialize agent memory\n",
        "memory = MemorySaver()\n",
        "\n",
        "\n",
        "# Define the agent's state\n",
        "class AgentState(TypedDict):\n",
        "    revision_number: int\n",
        "    max_revisions: int\n",
        "    search_count: int\n",
        "    max_searches: int\n",
        "    task: str\n",
        "    outline: str\n",
        "    queries: list\n",
        "    content: list\n",
        "    draft: str\n",
        "    critique: str\n",
        "    tool_calls: list"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "27b61a7e7ef6"
      },
      "source": [
        "### Initialize Gemini model\n",
        "\n",
        "Initialize the Gemini model from Vertex AI, specifying the model version and temperature settings.\n",
        "\n",
        "This sets up the core language model that will power your agent's actions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "06877aae6673"
      },
      "outputs": [],
      "source": [
        "model = ChatVertexAI(model=\"gemini-2.0-flash\", temperature=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d591fae74758"
      },
      "source": [
        "### Define search tools\n",
        "\n",
        "This section defines custom tools that will be used by your AI podcast agent to gather information from various sources. These tools act as interfaces to external services and provide access to relevant data for the podcast topic.\n",
        "\n",
        "Each tool is implemented as a Python function decorated with the `@tool` decorator from LangChain. This decorator makes it easy to integrate these functions into LangGraph workflows.\n",
        "\n",
        "The following search tools are defined:\n",
        "\n",
        "- **`search_arxiv`:** Retrieves research papers from arXiv based on a keyword query.\n",
        "- **`search_pubmed`:** Searches for information on PubMed, a database of biomedical literature.\n",
        "- **`search_wikipedia`:** Fetches information from Wikipedia based on a keyword query.\n",
        "\n",
        "Your LangGraph application will use these tool nodes to call the corresponding search functions and obtain information from these external sources. This allows the AI agent to learn about the podcast topic before generating the script."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0d27ed8a91c1"
      },
      "outputs": [],
      "source": [
        "@tool\n",
        "def search_arxiv(query: str) -> list[Document]:\n",
        "    \"\"\"Search for relevant publications on arXiv\"\"\"\n",
        "    retriever = ArxivRetriever(\n",
        "        load_max_docs=2,\n",
        "        get_full_documents=True,\n",
        "    )\n",
        "    docs = retriever.invoke(query)\n",
        "    if docs:\n",
        "        return docs\n",
        "    else:\n",
        "        return [\"No results found on arXiv\"]\n",
        "\n",
        "\n",
        "@tool\n",
        "def search_pubmed(query: str) -> list[Document]:\n",
        "    \"\"\"Search for information on PubMed\"\"\"\n",
        "    retriever = PubMedRetriever()\n",
        "    docs = retriever.invoke(query)\n",
        "    if docs:\n",
        "        return docs\n",
        "    else:\n",
        "        return [\"No results found on PubMed\"]\n",
        "\n",
        "\n",
        "@tool\n",
        "def search_wikipedia(query: str) -> list[Document]:\n",
        "    \"\"\"Search for information on Wikipedia\"\"\"\n",
        "    retriever = WikipediaRetriever()\n",
        "    docs = retriever.invoke(query)\n",
        "    if docs:\n",
        "        return docs\n",
        "    else:\n",
        "        return [\"No results found on Wikipedia\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "95fbea6aaa20"
      },
      "source": [
        "### Define workflow stages along with corresponding prompts and functions\n",
        "\n",
        "This section defines the different stages of the AI podcast agent's workflow and the corresponding prompt templates and node functions that drive each stage.\n",
        "\n",
        "Each stage represents a specific task in the podcast creation process, such as generating an outline, conducting research, writing the script, and providing a critique.\n",
        "\n",
        "For each stage, you'll define:\n",
        "\n",
        "- **Prompt Template:** A carefully crafted text prompt that instructs the Gemini language model on what to do at that stage. The prompt provides context, instructions, and any necessary input data.\n",
        "- **Node Function:** A Python function that encapsulates the logic for executing that stage. The function typically involves:\n",
        "    - Constructing the prompt with relevant information from the agent's state.\n",
        "    - Invoking the Gemini API with the prompt.\n",
        "    - Processing the model's response and updating the agent's state.\n",
        "\n",
        "These prompt templates and node functions are the building blocks of the LangGraph workflow that orchestrates the entire podcast creation process.\n",
        "\n",
        "#### Podcast outline node\n",
        "\n",
        "This node generates a structured outline for the podcast based on the user-provided topic:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4ce60bbc06e6"
      },
      "outputs": [],
      "source": [
        "OUTLINE_PROMPT = \"\"\"You are an expert writer tasked with writing a high level outline of an engaging 2-minute podcast.\n",
        "Write such an outline for the user provided topic. Give an outline of the podcast along with any\n",
        "relevant notes or instructions for the sections.\"\"\"\n",
        "\n",
        "\n",
        "# Generate an outline for the podcast based on the user-provided topic\n",
        "def podcast_outline_node(state: AgentState):\n",
        "    messages = [\n",
        "        SystemMessage(content=OUTLINE_PROMPT),\n",
        "        HumanMessage(content=state[\"task\"]),\n",
        "    ]\n",
        "    response = model.invoke(messages)\n",
        "    return {\"outline\": response.content}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5d7349c32d28"
      },
      "source": [
        "#### Research plan node\n",
        "\n",
        "This node formulates a search query based on the podcast topic and previous queries:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "87df19f53b95"
      },
      "outputs": [],
      "source": [
        "RESEARCH_PLAN_PROMPT = \"\"\"You are a researcher tasked with providing information that can\n",
        "be used when writing the following podcast. Generate one search query consisting of a few\n",
        "keywords that will be used to gather any relevant information. Do not output any information\n",
        "other than the query consisting of a few words.\n",
        "\n",
        "These were the past queries, do not repeat keywords from past queries in your newly generated query:\n",
        "---\n",
        "{queries}\"\"\"\n",
        "\n",
        "\n",
        "# Generates a search query based on the outline\n",
        "def research_plan_node(state: AgentState):\n",
        "    messages = [\n",
        "        SystemMessage(content=RESEARCH_PLAN_PROMPT.format(queries=state[\"queries\"])),\n",
        "        HumanMessage(content=state[\"task\"]),\n",
        "    ]\n",
        "    response = model.invoke(messages)\n",
        "    queries = state[\"queries\"]\n",
        "    if queries:\n",
        "        queries.append(response.content)\n",
        "    else:\n",
        "        queries = [response.content]\n",
        "    return {\"queries\": queries}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1719cf80233f"
      },
      "source": [
        "#### Research task node\n",
        "\n",
        "This node executes a search using the selected tool and query, retrieving relevant information for the podcast:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9bcbfe53b7d9"
      },
      "outputs": [],
      "source": [
        "RESEARCH_TASK_PROMPT = \"\"\"Use the available search tools and search queries to find information\n",
        "relevant to the podcast. Try searching different sources to obtain different articles. Try using\n",
        "different search tools than what was used previously so that you can obtain a broader range of\n",
        "information.\n",
        "\n",
        "These are the previous tool calls, so you can choose a different tool:\n",
        "---\n",
        "{tool_calls}\n",
        "---\n",
        "These are the previous search results, so you can aim for different sources and content:\n",
        "---\n",
        "{content}\"\"\"\n",
        "\n",
        "\n",
        "# Performs searches using tools\n",
        "def research_agent_node(state: AgentState):\n",
        "    tool_calls = state[\"tool_calls\"]\n",
        "    content = state[\"content\"]\n",
        "    queries = state[\"queries\"]\n",
        "    query = queries[-1]\n",
        "    messages = [\n",
        "        SystemMessage(\n",
        "            content=RESEARCH_TASK_PROMPT.format(tool_calls=tool_calls, content=content)\n",
        "        ),\n",
        "        HumanMessage(content=query),\n",
        "    ]\n",
        "\n",
        "    # Perform function calls\n",
        "    tools = [search_arxiv, search_pubmed, search_wikipedia]\n",
        "    model_with_tools = model.bind_tools(tools)\n",
        "    response_tool_calls = model_with_tools.invoke(messages)\n",
        "    if tool_calls:\n",
        "        tool_calls.append(response_tool_calls)\n",
        "    else:\n",
        "        tool_calls = [response_tool_calls]\n",
        "\n",
        "    # Defines a tool node based on search functions\n",
        "    tool_node = ToolNode(tools)\n",
        "    response = tool_node.invoke({\"messages\": [response_tool_calls]})\n",
        "\n",
        "    # Add the search results to the content list in the agent state\n",
        "    for message in response.get(\"messages\", []):\n",
        "        if isinstance(message, ToolMessage):\n",
        "            content.insert(0, message.content)\n",
        "\n",
        "    return {\n",
        "        \"content\": content,\n",
        "        \"tool_calls\": tool_calls,\n",
        "        \"search_count\": state[\"search_count\"] + 1,\n",
        "    }\n",
        "\n",
        "\n",
        "# Determine whether to continue research based on the number of searches performed\n",
        "def should_continue_tools(state: AgentState):\n",
        "    if state[\"search_count\"] > state[\"max_searches\"]:\n",
        "        return \"generate_script\"\n",
        "    else:\n",
        "        return \"research_plan\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9eaa6d0bff8d"
      },
      "source": [
        "#### Podcast writing node\n",
        "\n",
        "This node generates a draft podcast script using the outline and research results, aiming for an engaging and informative style:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1742523735e8"
      },
      "outputs": [],
      "source": [
        "WRITER_PROMPT = \"\"\"\n",
        "You are a writing assistant tasked with writing engaging 2-minute podcast scripts.\n",
        "\n",
        "- Generate the best podcast script possible for the user's request and the initial outline.\n",
        "- The script MUST strictly alternate lines between the two hosts, separating each host's line with a newline.\n",
        "- Add an intro phrase and outro phrase to start and end the podcast, and use a fun, random name for the podcast show.\n",
        "- Given a critique, respond with a revised version of your previous script.\n",
        "- Include lively back-and-forth chatter, reflections, and expressions of amazement between the hosts.\n",
        "- Cite at least THREE pieces of research throughout the script, choosing the most relevant research for each point.\n",
        "- DO NOT include ANY of the following:\n",
        "    - Speaker labels (e.g., \"Host 1:\", \"Host 2:\")\n",
        "    - Sound effect descriptions (e.g., \"[Sound of waves]\")\n",
        "    - Formatting instructions (e.g., \"(Emphasis)\", \"[Music fades in]\")\n",
        "    - Any other non-dialogue text.\n",
        "- Use this format for citations, including the month and year if available:\n",
        "    \"In [Month, Year], [Organization] found that...\"\n",
        "    \"Research from [Organization] in [Month, Year] showed that...\"\n",
        "    \"Back in [Month, Year], a study by [Organization] suggested that...\"\n",
        "---\n",
        "Utilize all of the following search results and context as needed:\n",
        "{content}\n",
        "---\n",
        "If this is a revision, the critique will be provided below:\n",
        "{critique}\"\"\"\n",
        "\n",
        "\n",
        "# Generates a draft of the script based on the content and outline\n",
        "def generate_script_node(state: AgentState):\n",
        "    messages = [\n",
        "        SystemMessage(\n",
        "            content=WRITER_PROMPT.format(\n",
        "                content=state[\"content\"], critique=state.get(\"critique\", \"\")\n",
        "            )\n",
        "        ),\n",
        "        HumanMessage(\n",
        "            content=f\"{state['task']}\\n\\nHere is my outline:\\n\\n{state['outline']}\"\n",
        "        ),\n",
        "    ]\n",
        "    response = model.invoke(messages)\n",
        "    return {\n",
        "        \"draft\": response.content,\n",
        "        \"search_count\": 0,  # Reset the search count for the next revision\n",
        "        \"revision_number\": state.get(\"revision_number\", 1) + 1,\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d00163e72e16"
      },
      "source": [
        "#### Podcast critique node\n",
        "\n",
        "This node provides feedback and suggestions for improvement on the generated podcast script:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "de70a68caa8d"
      },
      "outputs": [],
      "source": [
        "CRITIQUE_PROMPT = \"\"\"You are a producer grading a podcast script.\n",
        "Generate critique and recommendations for the user's submission.\n",
        "Provide detailed recommendations, including requests for conciseness, depth, style, etc.\"\"\"\n",
        "\n",
        "\n",
        "# Generates a critique with feedback on the draft podcast script\n",
        "def perform_critique_node(state: AgentState):\n",
        "    messages = [\n",
        "        SystemMessage(content=CRITIQUE_PROMPT),\n",
        "        HumanMessage(content=state[\"draft\"]),\n",
        "    ]\n",
        "    response = model.invoke(messages)\n",
        "    return {\"critique\": response.content}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "46aa739f9c1e"
      },
      "source": [
        "#### Research critique node\n",
        "\n",
        "This node generates a new search query based on the critique of the script, aiming to address weaknesses and find additional information:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "48682bcbb177"
      },
      "outputs": [],
      "source": [
        "RESEARCH_CRITIQUE_PROMPT = \"\"\"You are a writing assistant tasked with providing information that can\n",
        "be used when making any requested revisions (as outlined below).\n",
        "Generate one search query consisting of a few keywords that will be used to gather any relevant\n",
        "information. Do not output any information other than the query consisting of a few words.\n",
        "\n",
        "---\n",
        "\n",
        "These were the past queries, so you can vary the query that you generate:\n",
        "\n",
        "{queries}\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "# Generates a new search query based on the critique\n",
        "def research_critique_node(state: AgentState):\n",
        "    messages = [\n",
        "        SystemMessage(\n",
        "            content=RESEARCH_CRITIQUE_PROMPT.format(queries=state[\"queries\"])\n",
        "        ),\n",
        "        HumanMessage(content=state[\"critique\"]),\n",
        "    ]\n",
        "    response = model.invoke(messages)\n",
        "    queries = state.get(\"queries\", [])\n",
        "    if queries:\n",
        "        queries.append(response.content)\n",
        "    else:\n",
        "        queries = [response.content]\n",
        "    return {\"queries\": queries}\n",
        "\n",
        "\n",
        "# Decide whether to continue to the next revision or end the process\n",
        "def should_continue(state: AgentState):\n",
        "    if state[\"revision_number\"] > state[\"max_revisions\"]:\n",
        "        return END\n",
        "    return \"perform_critique\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aae38c0085e3"
      },
      "source": [
        "## Define and compile the LangGraph workflow\n",
        "\n",
        "This section defines the structure and flow of the AI podcast agent using LangGraph.\n",
        "\n",
        "The workflow is constructed as a graph with nodes representing each stage in the process (e.g., outlining, research, script generation) and edges defining the transitions between these stages.\n",
        "\n",
        "The workflow includes two main loops:\n",
        "\n",
        "- **Research Loop:**  This loop iteratively plans and executes searches until a specified number of searches are completed.\n",
        "- **Critique and Revision Loop:** This loop handles the script critique, additional research based on the critique, and script revision, repeating for a set number of revisions.\n",
        "\n",
        "The `workflow.compile()` function transforms this graph definition into an executable workflow, incorporating memory management to maintain the agent's state throughout the process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f7d04cda5f36"
      },
      "outputs": [],
      "source": [
        "# Initialize the state graph\n",
        "workflow = StateGraph(AgentState)\n",
        "\n",
        "# Define the nodes of the workflow, representing each stage of the process\n",
        "workflow.add_node(\"podcast_outline\", podcast_outline_node)\n",
        "workflow.add_node(\"research_plan\", research_plan_node)\n",
        "workflow.add_node(\"research_agent\", research_agent_node)\n",
        "workflow.add_node(\"generate_script\", generate_script_node)\n",
        "workflow.add_node(\"perform_critique\", perform_critique_node)\n",
        "workflow.add_node(\"research_critique\", research_critique_node)\n",
        "\n",
        "# Specify the starting node of the workflow\n",
        "workflow.set_entry_point(\"podcast_outline\")\n",
        "\n",
        "# Define the flow between node and stages\n",
        "workflow.add_edge(\"podcast_outline\", \"research_plan\")\n",
        "workflow.add_edge(\"research_plan\", \"research_agent\")\n",
        "workflow.add_edge(\"perform_critique\", \"research_critique\")\n",
        "workflow.add_edge(\"research_critique\", \"research_agent\")\n",
        "\n",
        "# Define conditional edges for the research loop\n",
        "workflow.add_conditional_edges(\n",
        "    \"research_agent\",\n",
        "    should_continue_tools,\n",
        "    {\"generate_script\": \"generate_script\", \"research_plan\": \"research_plan\"},\n",
        ")\n",
        "\n",
        "# Define conditional edges for the critique and revision loop\n",
        "workflow.add_conditional_edges(\n",
        "    \"generate_script\",\n",
        "    should_continue,\n",
        "    {END: END, \"perform_critique\": \"perform_critique\"},\n",
        ")\n",
        "\n",
        "# Compile the workflow graph, enabling memory to track agent state\n",
        "graph = workflow.compile(checkpointer=memory)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "18154fa6d8b4"
      },
      "source": [
        "### Visualize the workflow\n",
        "\n",
        "This cell visualizes the compiled LangGraph workflow as a [Mermaid diagram](https://mermaid.js.org/).\n",
        "\n",
        "The diagram provides a clear and intuitive representation of the workflow's structure, showing the nodes, edges, and the flow of execution.\n",
        "\n",
        "This visualization helps to understand the overall process and the interactions between different stages of the AI podcast agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f97fe13cc0cf"
      },
      "outputs": [],
      "source": [
        "# Display a Mermaid diagram of the workflow\n",
        "Image(graph.get_graph().draw_mermaid_png())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3076e0cd0c4e"
      },
      "source": [
        "### Define the podcast topic\n",
        "\n",
        "This cell defines the topic of the podcast that the AI agent will create.\n",
        "\n",
        "The topic is assigned to the variable `PODCAST_TOPIC`. Feel free to modify this variable to explore different podcast topics! "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6627153c6715"
      },
      "outputs": [],
      "source": [
        "PODCAST_TOPIC = \"Explore the use of bio-inspired fluid dynamics in the design of underwater robots and vehicles\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a1df657ce053"
      },
      "source": [
        "### Run the AI podcast agent\n",
        "\n",
        "This cell executes the compiled LangGraph workflow, running the AI podcast agent to generate the podcast script. \n",
        "\n",
        "The code performs these actions:\n",
        "\n",
        "- **Clean agent helper function:**  This function prepares the agent's output for printing by removing unnecessary characters and formatting\n",
        "- **Thread Configuration:**  A thread configuration is defined to ensure a unique history for this workflow execution\n",
        "- **Workflow Execution:** The `graph.stream()` method runs the workflow, iterating through each stage and updating the agent's state\n",
        "- **Output Display:** The code prints the results of each stage, including the agent's actions and generated output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "338377bc8c25"
      },
      "outputs": [],
      "source": [
        "# Function to clean and format agent output for display\n",
        "\n",
        "\n",
        "def clean_agent_result(data):\n",
        "    agent_result = str(data)\n",
        "    agent_result = re.sub(\n",
        "        r\"[^\\x00-\\x7F]+\", \" \", agent_result\n",
        "    )  # Remove non-ASCII characters\n",
        "    agent_result = re.sub(r\"\\\\\\\\n\", \"\\n\", agent_result)  # Replace escaped newlines\n",
        "    agent_result = re.sub(r\"\\\\n\", \"\", agent_result)  # Replace newlines\n",
        "    agent_result = re.sub(r\"\\\\'\", \"'\", agent_result)  # Replace escaped single quotes\n",
        "    return agent_result\n",
        "\n",
        "\n",
        "# Thread ID for unique history in workflow execution\n",
        "thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
        "\n",
        "# Run the LangGraph workflow, passing the initial state and thread configuration\n",
        "for state in graph.stream(\n",
        "    {\n",
        "        \"task\": PODCAST_TOPIC,\n",
        "        \"revision_number\": 1,  # Current revision number\n",
        "        \"max_revisions\": 2,  # Maximum number of revisions allowed\n",
        "        \"search_count\": 0,  # Current search number\n",
        "        \"max_searches\": 3,  # Maximum number of searches allowed per revision\n",
        "        \"content\": [],\n",
        "        \"queries\": [],\n",
        "        \"tool_calls\": [],\n",
        "    },\n",
        "    thread,\n",
        "):\n",
        "    # Print a snippet of the results of each workflow stage\n",
        "    for k, v in state.items():\n",
        "        print(f\"Agent Node: {k}\\n\")\n",
        "        print(\"Agent Result:\")\n",
        "        print(clean_agent_result(v)[:1000])\n",
        "    print(\"\\n====================\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cd8d951762fb"
      },
      "source": [
        "### Parse and display the final podcast script\n",
        "\n",
        "This section extracts and prepares the final podcast script generated by the AI agent.\n",
        "\n",
        "It displays the script for review, where each string in the list will be narrated by a different text-to-speech voice."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "599397cab03c"
      },
      "outputs": [],
      "source": [
        "podcast_script = state[\"generate_script\"][\"draft\"]\n",
        "parsed_script = [\n",
        "    text for text in (line.strip() for line in podcast_script.splitlines()) if text\n",
        "]\n",
        "parsed_script"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e075edfe820d"
      },
      "source": [
        "### Generate audio for the podcast\n",
        "\n",
        "This cell generates audio for each line of the parsed podcast script using Google Cloud's Text-to-Speech API.\n",
        "\n",
        "It creates separate audio files for each line, alternating between two different voices to simulate a conversation between two podcast hosts.\n",
        "\n",
        "The code:\n",
        "\n",
        "1. **Initializes the Text-to-Speech Client:**  Sets up the connection to the API.\n",
        "2. **Defines Audio Configuration:** Specifies the desired output audio format (MP3).\n",
        "3. **Iterates through Script Lines:** Generates audio for each line, alternating voices.\n",
        "4. **Saves Audio Files:**  Writes the generated audio to separate MP3 files."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c74e3badfd35"
      },
      "outputs": [],
      "source": [
        "# Instantiates a client\n",
        "client = texttospeech.TextToSpeechClient()\n",
        "\n",
        "# Select the type of audio file you want returned\n",
        "audio_config = texttospeech.AudioConfig(audio_encoding=texttospeech.AudioEncoding.MP3)\n",
        "\n",
        "audio_files = []\n",
        "for count, line in enumerate(parsed_script):\n",
        "    # Set the text input to be synthesized\n",
        "    synthesis_input = texttospeech.SynthesisInput(text=line)\n",
        "\n",
        "    # Choose the voice for the current line, alternating between hosts\n",
        "    if count % 2 == 0:\n",
        "        voice_name = \"en-US-Chirp3-HD-Aoede\"\n",
        "    elif count % 2 == 1:\n",
        "        voice_name = \"en-US-Chirp3-HD-Puck\"\n",
        "\n",
        "    # Configure voice parameters: language and voice name\n",
        "    voice = texttospeech.VoiceSelectionParams(\n",
        "        language_code=\"en-US\",\n",
        "        name=voice_name,\n",
        "    )\n",
        "\n",
        "    # Generate audio using the Text-to-Speech API\n",
        "    response = client.synthesize_speech(\n",
        "        input=synthesis_input, voice=voice, audio_config=audio_config\n",
        "    )\n",
        "\n",
        "    # Save the generated audio to an MP3 file\n",
        "    filename = f\"part-{str(count)}.mp3\"\n",
        "    audio_files.append(filename)\n",
        "    with open(filename, \"wb\") as out:\n",
        "        out.write(response.audio_content)\n",
        "        print(f\"Audio content written to file {filename}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ea674bbf2628"
      },
      "source": [
        "### Combine audio files and generate final podcast\n",
        "\n",
        "This cell combines the individual audio files generated in the previous step into a single podcast file.\n",
        "\n",
        "It also adds brief silences between each line for better listening experience.\n",
        "\n",
        "The final podcast is saved as `gemini-podcast.mp3`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9a4e93adc415"
      },
      "outputs": [],
      "source": [
        "# Initialize audio segment\n",
        "full_audio = AudioSegment.silent(duration=200)\n",
        "\n",
        "# Concatenate audio segments with silence in between\n",
        "for file in audio_files:\n",
        "    sound = AudioSegment.from_mp3(file)\n",
        "    silence = AudioSegment.silent(duration=200)\n",
        "    full_audio += sound + silence\n",
        "    os.remove(file)  # Remove the individual part files after combining\n",
        "\n",
        "# Save the final audio output to a file\n",
        "podcast_filename = \"gemini-podcast.mp3\"\n",
        "full_audio.export(podcast_filename)\n",
        "print(f\"Podcast content written to file {podcast_filename}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2be9b47ae8a3"
      },
      "source": [
        "### Listen to your AI-generated podcast!\n",
        "\n",
        "This cell plays the final podcast generated by the AI agent.\n",
        "\n",
        "The `Audio` object from `IPython.display` is used to embed the audio player directly into the notebook. The podcast will start playing automatically.\n",
        "\n",
        "Enjoy your AI-created podcast!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2f50dcfe1651"
      },
      "outputs": [],
      "source": [
        "Audio(filename=podcast_filename, rate=32000, autoplay=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "df34f0351197"
      },
      "source": [
        "## Conclusion: Building AI-powered podcast agents\n",
        "\n",
        "This notebook showcases the exciting potential of using AI to automate the podcast creation process. By combining the power of the Gemini API with the flexibility of LangGraph, you built an intelligent agent capable of:\n",
        "\n",
        "- **Generating Podcast Outlines:** Structuring the flow and content of the podcast.\n",
        "- **Conducting Research:**  Gathering information from various sources like arXiv, PubMed, and Wikipedia.\n",
        "- **Writing Engaging Scripts:** Crafting podcast scripts with natural-sounding dialogue, citations, and a conversational style.\n",
        "- **Critiquing and Revising:** Providing feedback on the script and iteratively refining it.\n",
        "- **Generating Audio:**  Using text-to-speech technology to create the final podcast audio.\n",
        "\n",
        "This is just a starting point! You can customize this workflow further by:\n",
        "\n",
        "- **Adding New Research Tools:**  Integrate additional sources of information relevant to your podcast topics.\n",
        "- **Experimenting with Prompts:** Refine the prompts to guide the AI agent towards your desired style and content.\n",
        "- **Exploring Different Voices:**  Use a wider range of voices for the podcast hosts to create unique and engaging listening experiences.\n",
        "\n",
        "The possibilities are endless! As AI technology continues to advance, you can expect even more creative and innovative applications in podcasting and other content creation domains.\n",
        "\n",
        "You can learn more about [LangGraph](https://langchain-ai.github.io/langgraph/), the [Gemini API in Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models), or the [chat model provider for Vertex AI in LangChain](https://python.langchain.com/docs/integrations/chat/google_vertex_ai_palm/) in their respective documentation pages."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "langgraph_gemini_podcast.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
