1_agentic-design-ptn/02_plan-and-execute/LangGraph/01.1_plan-and-execute.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5c95e9bf",
   "metadata": {},
   "source": [
    "# Plan-and-Execute\n",
    "---\n",
    "\n",
    "### What is Plan-and-Execute?\n",
    "The Plan-and-Execute framework is a strategy for retrieval-augmented generation (RAG) that divides complex reasoning tasks into two distinct phases: planning and execution. While traditional ReAct agents think one step at a time, plan-and-execute emphasizes explicit, long-term planning.\n",
    "\n",
    "- **Planning Phase**: The model generates a high-level plan or structured outline that serves as a roadmap for solving the task. This phase ensures that the execution is systematic and adheres to the task's requirements.\n",
    "- **Execution Phase**: Based on the generated plan, the model retrieves relevant information and executes the outlined steps to provide a detailed and coherent response.\n",
    "\n",
    "This separation aims to address limitations in RAG systems that attempt to perform reasoning and generation in a single step, often leading to logical errors or inefficiency in handling complex tasks.\n",
    "\n",
    "### Key Advantages\n",
    "- **Improved Task Decomposition**: By explicitly separating planning and execution, the framework enables better handling of complex, multi-step reasoning tasks, ensuring systematic progress towards the solution.\n",
    "- **Higher Accuracy and Coherence**: The planning phase acts as a guide, reducing the chances of errors and improving the logical coherence of the responses generated during execution.\n",
    "\n",
    "**Reference**\n",
    "- [ReAct paper](https://arxiv.org/abs/2210.03629)\n",
    "- [Plan-and-Solve paper](https://arxiv.org/abs/2305.04091)\n",
    "- [Baby-AGI project](https://github.com/yoheinakajima/babyagi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb1ad59b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from azure_genai_utils.tracer import get_langchain_api_key, set_langsmith\n",
    "\n",
    "load_dotenv(override=True)\n",
    "\n",
    "# If you want to trace your RAG API calls, please set the tracing=True. You need to have a valid Langchain API key.\n",
    "langchain_key, has_langchain_key = get_langchain_api_key()\n",
    "set_langsmith(\"[RAG Innv Lab] 1_Agentic-Design-Pattern\", tracing=False)\n",
    "\n",
    "azure_openai_chat_deployment_name = os.getenv(\"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME\")\n",
    "azure_openai_embedding_deployment_name = os.getenv(\"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a56c472",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LANGUAGE = \"English\"\n",
    "# LOCALE = \"en-US\"\n",
    "LANGUAGE = \"Korean\"\n",
    "LOCALE = \"ko-KR\"\n",
    "\n",
    "language_prompt = f\" Answer in {LANGUAGE}.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0856a092",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 🧪 Step 1. Test and Construct each module\n",
    "---\n",
    "\n",
    "Before building the entire the graph pipeline, we will test and construct each module separately.\n",
    "\n",
    "### Web Search Tool\n",
    "\n",
    "Web search tool is used to enhance the context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16767612",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure_genai_utils.tools import BingSearch\n",
    "\n",
    "WEB_SEARCH_FORMAT_OUTPUT = False\n",
    "\n",
    "web_search_tool = BingSearch(\n",
    "    max_results=2,\n",
    "    locale=LOCALE,\n",
    "    include_news=False,\n",
    "    include_entity=False,\n",
    "    format_output=WEB_SEARCH_FORMAT_OUTPUT,\n",
    ")\n",
    "tools = [web_search_tool]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a04fa1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Microsoft AutoGen\"\n",
    "results = web_search_tool.invoke({\"query\": question})\n",
    "print(results[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "348ac462",
   "metadata": {},
   "source": [
    "### Define your LLM\n",
    "\n",
    "This hands-on only uses the `gpt-4o`, but you can utilize multiple models in the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa5217d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import AzureChatOpenAI\n",
    "\n",
    "model_name = azure_openai_chat_deployment_name\n",
    "llm = AzureChatOpenAI(model=model_name, temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d230b5c",
   "metadata": {},
   "source": [
    "### Execution Agent\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f70cf9fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.prebuilt import create_react_agent\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            f\"You are a helpful assistant. {language_prompt}\",\n",
    "        ),\n",
    "        (\"human\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Create the ReAct agent\n",
    "agent_executor = create_react_agent(llm, tools, prompt=prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab1307f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_executor.invoke({\"messages\": [(\"user\", \"Where is Seoul?\")]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e125cecc",
   "metadata": {},
   "source": [
    "### Plan Phase\n",
    "Now let's consider how to create a **planning phase**, where we use function calling to create a plan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f84bd07",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from typing import Annotated, List\n",
    "\n",
    "\n",
    "class Plan(BaseModel):\n",
    "    \"\"\"Sorted steps to execute the plan\"\"\"\n",
    "\n",
    "    steps: Annotated[List[str], \"Different steps to follow, should be in sorted order\"]\n",
    "\n",
    "\n",
    "planner_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            f\"\"\"For the given objective, come up with a simple step by step plan. \\\n",
    "This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \\\n",
    "The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.\n",
    "{language_prompt}\"\"\",\n",
    "        ),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "planner = planner_prompt | llm.with_structured_output(Plan)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c80a8f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "planner.invoke(\n",
    "    {\"messages\": [((\"user\", \"What is the benefit of using Microsoft AutoGen?\"))]}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87af450a",
   "metadata": {},
   "source": [
    "### Re-plan phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b2351a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Union\n",
    "\n",
    "\n",
    "class Response(BaseModel):\n",
    "    \"\"\"Response to user.\"\"\"\n",
    "\n",
    "    response: str\n",
    "\n",
    "\n",
    "class Act(BaseModel):\n",
    "    \"\"\"Action to perform.\"\"\"\n",
    "\n",
    "    # Use \"Response\" if you want to respond to user, use \"Plan\" if you need to further use tools to get the answer.\n",
    "    action: Union[Response, Plan] = Field(\n",
    "        description=\"Action to perform. If you want to respond to user, use 'Response'. \"\n",
    "        \"If you need to further use tools to get the answer, use 'Plan'.\"\n",
    "    )\n",
    "\n",
    "\n",
    "replanner_prompt = ChatPromptTemplate.from_template(\n",
    "    \"\"\"For the given objective, come up with a simple step by step plan. \\\n",
    "This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \\\n",
    "The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.\n",
    "\n",
    "Your objective was this:\n",
    "{input}\n",
    "\n",
    "Your original plan was this:\n",
    "{plan}\n",
    "\n",
    "You have currently done the follow steps:\n",
    "{past_steps}\n",
    "\n",
    "Update your plan accordingly. If no more steps are needed and you can return to the user, then respond with that. Otherwise, fill out the plan. Only add steps to the plan that still NEED to be done. Do not return previously done steps as part of the plan.\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "replanner = replanner_prompt | llm.with_structured_output(Act)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc0058df",
   "metadata": {},
   "outputs": [],
   "source": [
    "replanner.invoke(\n",
    "    {\n",
    "        \"input\": \"What is the benifit of using Microsoft AutoGen?\",\n",
    "        \"plan\": \"Step 1: Search for the benefits of Microsoft AutoGen.\",\n",
    "        \"past_steps\": \"\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a9c06ee",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 🧪 Step 2. Define the Graph\n",
    "---\n",
    "\n",
    "### State Definition\n",
    "\n",
    "- `input`: User input\n",
    "- `plan`: Current plan\n",
    "- `past_steps`: past steps and results\n",
    "- `response`: final response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27c44393",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated, List, Tuple\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "\n",
    "class PlanExecute(TypedDict):\n",
    "    input: Annotated[str, \"User's input\"]\n",
    "    plan: Annotated[List[str], \"Current plan\"]\n",
    "    past_steps: Annotated[List[Tuple], operator.add]\n",
    "    response: Annotated[str, \"Final response\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b88a7250",
   "metadata": {},
   "source": [
    "### Define Nodes\n",
    "\n",
    "We will define the following nodes in the graph:\n",
    "\n",
    "- `plan_step`: Plan the steps to execute the given task.\n",
    "- `execute_step`: Execute the given task using Agent executor.\n",
    "- `replan_step`: Replan from previous steps or get final response.\n",
    "- `should_end`: Check if the agent should end execution or continue.\n",
    "- `generate_final_report`: Generate a final report based on the execution results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df12d945",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "\n",
    "def plan_step(state: PlanExecute):\n",
    "    \"\"\"Plan the steps to execute the given task.\"\"\"\n",
    "    plan = planner.invoke({\"messages\": [(\"user\", state[\"input\"])]})\n",
    "    return {\"plan\": plan.steps}\n",
    "\n",
    "\n",
    "def execute_step(state: PlanExecute):\n",
    "    \"\"\"Execute the given task using Agent executor.\"\"\"\n",
    "    plan = state[\"plan\"]\n",
    "    # Each step is numbered\n",
    "    plan_str = \"\\n\".join(f\"{i+1}. {step}\" for i, step in enumerate(plan))\n",
    "    task = plan[0]\n",
    "    task_formatted = f\"\"\"For the following plan:\n",
    "{plan_str}\\n\\nYou are tasked with executing [step 1. {task}].\"\"\"\n",
    "    agent_response = agent_executor.invoke({\"messages\": [(\"user\", task_formatted)]})\n",
    "\n",
    "    return {\n",
    "        \"past_steps\": [(task, agent_response[\"messages\"][-1].content)],\n",
    "    }\n",
    "\n",
    "\n",
    "def replan_step(state: PlanExecute):\n",
    "    \"\"\"Replan from previous steps or get final response.\"\"\"\n",
    "    output = replanner.invoke(state)\n",
    "\n",
    "    if isinstance(output.action, Response):\n",
    "        return {\"response\": output.action.response}\n",
    "    else:\n",
    "        next_plan = output.action.steps\n",
    "        if len(next_plan) == 0:\n",
    "            return {\"response\": \"No more steps needed.\"}\n",
    "        else:\n",
    "            return {\"plan\": next_plan}\n",
    "\n",
    "\n",
    "def should_end(state: PlanExecute):\n",
    "    \"\"\"Check if the agent should end execution or continue.\"\"\"\n",
    "    if \"response\" in state and state[\"response\"]:\n",
    "        return \"final_report\"\n",
    "    else:\n",
    "        return \"execute\"\n",
    "\n",
    "\n",
    "final_report_prompt = ChatPromptTemplate.from_template(\n",
    "    \"\"\"You are given the objective and the previously done steps. Your task is to generate a final report in markdown format.\n",
    "Final report should be written in professional tone.\n",
    "\n",
    "Your objective was this:\n",
    "\n",
    "{input}\n",
    "\n",
    "Your previously done steps(question and answer pairs):\n",
    "\n",
    "{past_steps}\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "final_report = final_report_prompt | llm | StrOutputParser()\n",
    "\n",
    "\n",
    "def generate_final_report(state: PlanExecute):\n",
    "    \"\"\"Generate a final report based on the execution results.\"\"\"\n",
    "    past_steps = \"\\n\\n\".join(\n",
    "        [\n",
    "            f\"Question: {past_step[0]}\\n\\nAnswer: {past_step[1]}\\n\\n####\"\n",
    "            for past_step in state[\"past_steps\"]\n",
    "        ]\n",
    "    )\n",
    "    response = final_report.invoke({\"input\": state[\"input\"], \"past_steps\": past_steps})\n",
    "    return {\"response\": response}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff7ebe60",
   "metadata": {},
   "source": [
    "### Construct the Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16abdcf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "workflow = StateGraph(PlanExecute)\n",
    "\n",
    "# Node definitions\n",
    "workflow.add_node(\"planner\", plan_step)\n",
    "workflow.add_node(\"execute\", execute_step)\n",
    "workflow.add_node(\"replan\", replan_step)\n",
    "workflow.add_node(\"final_report\", generate_final_report)\n",
    "\n",
    "# Edge connections\n",
    "workflow.add_edge(START, \"planner\")\n",
    "workflow.add_edge(\"planner\", \"execute\")\n",
    "workflow.add_edge(\"execute\", \"replan\")\n",
    "workflow.add_edge(\"final_report\", END)\n",
    "workflow.add_conditional_edges(\n",
    "    \"replan\",\n",
    "    should_end,\n",
    "    {\"execute\": \"execute\", \"final_report\": \"final_report\"},\n",
    ")\n",
    "\n",
    "# Compile the workflow\n",
    "app = workflow.compile(checkpointer=MemorySaver())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd850b8c",
   "metadata": {},
   "source": [
    "### Visualize the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ee7d379",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure_genai_utils.graphs import visualize_langgraph\n",
    "\n",
    "visualize_langgraph(app, xray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6d266cb",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 🧪 Step 3. Execute the Graph\n",
    "---\n",
    "\n",
    "### Execute the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6323963b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure_genai_utils.messages import invoke_graph, random_uuid\n",
    "from langchain_core.runnables import RunnableConfig\n",
    "\n",
    "config = RunnableConfig(recursion_limit=50, configurable={\"thread_id\": random_uuid()})\n",
    "\n",
    "inputs = {\n",
    "    \"input\": f\"Explain the main differences between Microsoft AutoGen and LangGraph. {language_prompt}\"\n",
    "}\n",
    "\n",
    "invoke_graph(app, inputs, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0e9c85e",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "611af101",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown\n",
    "\n",
    "snapshot = app.get_state(config).values\n",
    "Markdown(snapshot[\"response\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "505af72e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_agent",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}