{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "457cdc67-1893-4653-8b0c-b185a5947e74",
   "metadata": {},
   "source": [
    "# How to migrate from legacy LangChain agents to LangGraph\n",
    "\n",
    "Here we focus on how to move from legacy LangChain agents to LangGraph agents.\n",
    "LangChain agents (the [AgentExecutor](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor) in particular) have multiple configuration parameters.\n",
    "In this notebook we will show how those parameters map to the LangGraph [react agent executor](https://langchain-ai.github.io/langgraph/reference/prebuilt/#create_react_agent).\n",
    "\n",
    "#### Prerequisites\n",
    "\n",
    "This how-to guide uses OpenAI as the LLM. Install the dependencies to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "662fac50",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph langchain langchain-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e50635c-1671-46e6-be65-ce95f8167c2f",
   "metadata": {},
   "source": [
    "## Basic Usage\n",
    "\n",
    "For basic creation and usage of a tool-calling ReAct-style agent, the functionality is the same. First, let's define a model and tool(s), then we'll use those to create an agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1e425fea-2796-4b99-bee6-9a6ffe73f756",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o\")\n",
    "\n",
    "\n",
    "@tool\n",
    "def magic_function(input: int) -> int:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    return input + 2\n",
    "\n",
    "\n",
    "tools = [magic_function]\n",
    "\n",
    "\n",
    "query = \"what is the value of magic_function(3)?\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af002033-fe51-4d14-b47c-3e9b483c8395",
   "metadata": {},
   "source": [
    "For the LangChain [AgentExecutor](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor), we define a prompt with a placeholder for the agent's scratchpad. The agent can be invoked as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "03ea357c-9c36-4464-b2cc-27bd150e1554",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input': 'what is the value of magic_function(3)?',\n",
       " 'output': 'The value of `magic_function(3)` is 5.'}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.agents import AgentExecutor, create_tool_calling_agent\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "        # Placeholders fill up a **list** of messages\n",
    "        (\"placeholder\", \"{agent_scratchpad}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "agent = create_tool_calling_agent(model, tools, prompt)\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools)\n",
    "\n",
    "agent_executor.invoke({\"input\": query})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94205f3b-fd2b-4fd7-af69-0a3fc313dc88",
   "metadata": {},
   "source": [
    "LangGraph's [react agent executor](https://langchain-ai.github.io/langgraph/reference/prebuilt/#create_react_agent) manages a state that is defined by a list of messages. It will continue to process the list until there are no tool calls in the agent's output. To kick it off, we input a list of messages. The output will contain the entire state of the graph-- in this case, the conversation history.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "53a3737a-d167-4255-89bf-20ac37f89a3e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input': 'what is the value of magic_function(3)?',\n",
       " 'output': 'The value of `magic_function(3)` is 5.'}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "app = create_react_agent(model, tools)\n",
    "\n",
    "\n",
    "messages = app.invoke({\"messages\": [(\"human\", query)]})\n",
    "{\n",
    "    \"input\": query,\n",
    "    \"output\": messages[\"messages\"][-1].content,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "74ecebe3-512e-409c-a661-bdd5b0a2b782",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input': 'Pardon?',\n",
       " 'output': 'The result of applying `magic_function` to the input 3 is 5.'}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "message_history = messages[\"messages\"]\n",
    "\n",
    "new_query = \"Pardon?\"\n",
    "\n",
    "messages = app.invoke({\"messages\": message_history + [(\"human\", new_query)]})\n",
    "{\n",
    "    \"input\": new_query,\n",
    "    \"output\": messages[\"messages\"][-1].content,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4466a4d-e55e-4ece-bee8-2269a0b5677b",
   "metadata": {},
   "source": [
    "## Prompt Templates\n",
    "\n",
    "With legacy LangChain agents you have to pass in a prompt template. You can use this to control the agent.\n",
    "\n",
    "With LangGraph [react agent executor](https://langchain-ai.github.io/langgraph/reference/prebuilt/#create_react_agent), by default there is no prompt. You can achieve similar control over the agent in a few ways:\n",
    "\n",
    "1. Pass in a system message as input\n",
    "2. Initialize the agent with a system message\n",
    "3. Initialize the agent with a function to transform messages before passing to the model.\n",
    "\n",
    "Let's take a look at all of these below. We will pass in custom instructions to get the agent to respond in Spanish.\n",
    "\n",
    "First up, using AgentExecutor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a9a11ccd-75e2-4c11-844d-a34870b0ff91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input': 'what is the value of magic_function(3)?',\n",
       " 'output': 'El valor de `magic_function(3)` es 5.'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant. Respond only in Spanish.\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "        # Placeholders fill up a **list** of messages\n",
    "        (\"placeholder\", \"{agent_scratchpad}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "agent = create_tool_calling_agent(model, tools, prompt)\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools)\n",
    "\n",
    "agent_executor.invoke({\"input\": query})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd5f5500-5ae4-4000-a9fd-8c5a2cc6404d",
   "metadata": {},
   "source": [
    "Now, let's pass a custom system message to [react agent executor](https://langchain-ai.github.io/langgraph/reference/prebuilt/#create_react_agent). This can either be a string or a LangChain SystemMessage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a9486805-676a-4d19-a5c4-08b41b172989",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import SystemMessage\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "system_message = \"You are a helpful assistant. Respond only in Spanish.\"\n",
    "# This could also be a SystemMessage object\n",
    "# system_message = SystemMessage(content=\"You are a helpful assistant. Respond only in Spanish.\")\n",
    "\n",
    "app = create_react_agent(model, tools, messages_modifier=system_message)\n",
    "\n",
    "\n",
    "messages = app.invoke({\"messages\": [(\"user\", query)]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc6059fd-0df7-4b6f-a84c-b5874e983638",
   "metadata": {},
   "source": [
    "We can also pass in an arbitrary function. This function should take in a list of messages and output a list of messages.\n",
    "We can do all types of arbitrary formatting of messages here. In this cases, let's just add a SystemMessage to the start of the list of messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d369ab45-0c82-45f4-9d3e-8efb8dd47e2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input': 'what is the value of magic_function(3)?', 'output': 'El valor de magic_function(3) es 5. ¡Pandamonium!'}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import AnyMessage\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant. Respond only in Spanish.\"),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "def _modify_messages(messages: list[AnyMessage]):\n",
    "    return prompt.invoke({\"messages\": messages}).to_messages() + [\n",
    "        (\"user\", \"Also say 'Pandamonium!' after the answer.\")\n",
    "    ]\n",
    "\n",
    "\n",
    "app = create_react_agent(model, tools, messages_modifier=_modify_messages)\n",
    "\n",
    "\n",
    "messages = app.invoke({\"messages\": [(\"human\", query)]})\n",
    "print(\n",
    "    {\n",
    "        \"input\": query,\n",
    "        \"output\": messages[\"messages\"][-1].content,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68df3a09",
   "metadata": {},
   "source": [
    "## Memory\n",
    "\n",
    "With LangChain's [AgentExecutor](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor.iter), you could add chat [Memory](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor.memory) so it can engage in a multi-turn conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1fb52a2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi Polly! The output of the magic function for the input 3 is 5.\n",
      "---\n",
      "Yes, I remember your name, Polly! How can I assist you further?\n",
      "---\n",
      "The output of the magic function for the input 3 is 5.\n"
     ]
    }
   ],
   "source": [
    "from langchain.agents import AgentExecutor, create_tool_calling_agent\n",
    "from langchain_community.chat_message_histories import ChatMessageHistory\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o\")\n",
    "memory = ChatMessageHistory(session_id=\"test-session\")\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant.\"),\n",
    "        # First put the history\n",
    "        (\"placeholder\", \"{chat_history}\"),\n",
    "        # Then the new input\n",
    "        (\"human\", \"{input}\"),\n",
    "        # Finally the scratchpad\n",
    "        (\"placeholder\", \"{agent_scratchpad}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "@tool\n",
    "def magic_function(input: int) -> int:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    return input + 2\n",
    "\n",
    "\n",
    "tools = [magic_function]\n",
    "\n",
    "\n",
    "agent = create_tool_calling_agent(model, tools, prompt)\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools)\n",
    "\n",
    "agent_with_chat_history = RunnableWithMessageHistory(\n",
    "    agent_executor,\n",
    "    # This is needed because in most real world scenarios, a session id is needed\n",
    "    # It isn't really used here because we are using a simple in memory ChatMessageHistory\n",
    "    lambda session_id: memory,\n",
    "    input_messages_key=\"input\",\n",
    "    history_messages_key=\"chat_history\",\n",
    ")\n",
    "\n",
    "config = {\"configurable\": {\"session_id\": \"test-session\"}}\n",
    "print(\n",
    "    agent_with_chat_history.invoke(\n",
    "        {\"input\": \"Hi, I'm polly! What's the output of magic_function of 3?\"}, config\n",
    "    )[\"output\"]\n",
    ")\n",
    "print(\"---\")\n",
    "print(agent_with_chat_history.invoke({\"input\": \"Remember my name?\"}, config)[\"output\"])\n",
    "print(\"---\")\n",
    "print(\n",
    "    agent_with_chat_history.invoke({\"input\": \"what was that output again?\"}, config)[\n",
    "        \"output\"\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2a5a32f",
   "metadata": {},
   "source": [
    "#### In LangGraph\n",
    "\n",
    "Memory is just [persistence](https://langchain-ai.github.io/langgraph/how-tos/persistence/), aka [checkpointing](https://langchain-ai.github.io/langgraph/reference/checkpoints/).\n",
    "\n",
    "Add a `checkpointer` to the agent and you get chat memory for free."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "035e1253",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi Polly! The output of the magic_function for the input 3 is 5.\n",
      "---\n",
      "Yes, your name is Polly!\n",
      "---\n",
      "The output of the magic_function for the input 3 was 5.\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import SystemMessage\n",
    "from langgraph.checkpoint import MemorySaver  # an in-memory checkpointer\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "system_message = \"You are a helpful assistant.\"\n",
    "# This could also be a SystemMessage object\n",
    "# system_message = SystemMessage(content=\"You are a helpful assistant. Respond only in Spanish.\")\n",
    "\n",
    "memory = MemorySaver()\n",
    "app = create_react_agent(\n",
    "    model, tools, messages_modifier=system_message, checkpointer=memory\n",
    ")\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"test-thread\"}}\n",
    "print(\n",
    "    app.invoke(\n",
    "        {\n",
    "            \"messages\": [\n",
    "                (\"user\", \"Hi, I'm polly! What's the output of magic_function of 3?\")\n",
    "            ]\n",
    "        },\n",
    "        config,\n",
    "    )[\"messages\"][-1].content\n",
    ")\n",
    "print(\"---\")\n",
    "print(\n",
    "    app.invoke({\"messages\": [(\"user\", \"Remember my name?\")]}, config)[\"messages\"][\n",
    "        -1\n",
    "    ].content\n",
    ")\n",
    "print(\"---\")\n",
    "print(\n",
    "    app.invoke({\"messages\": [(\"user\", \"what was that output again?\")]}, config)[\n",
    "        \"messages\"\n",
    "    ][-1].content\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7cf24a8",
   "metadata": {},
   "source": [
    "## Iterating through steps\n",
    "\n",
    "With LangChain's [AgentExecutor](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor.iter), you could iterate over the steps using the [stream](https://api.python.langchain.com/en/latest/runnables/langchain_core.runnables.base.Runnable.html#langchain_core.runnables.base.Runnable.stream) (or async `astream`) methods or the [iter](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor.iter) method. LangGraph supports stepwise iteration using [stream](https://api.python.langchain.com/en/latest/runnables/langchain_core.runnables.base.Runnable.html#langchain_core.runnables.base.Runnable.stream) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d640feb3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'actions': [ToolAgentAction(tool='magic_function', tool_input={'input': 3}, log=\"\\nInvoking: `magic_function` with `{'input': 3}`\\n\\n\\n\", message_log=[AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_q9MgGFjqJbV2xSUX93WqxmOt', 'function': {'arguments': '{\"input\":3}', 'name': 'magic_function'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-c68fd76f-a3c3-4c3c-bfd7-748c171ed4b8', tool_calls=[{'name': 'magic_function', 'args': {'input': 3}, 'id': 'call_q9MgGFjqJbV2xSUX93WqxmOt'}], tool_call_chunks=[{'name': 'magic_function', 'args': '{\"input\":3}', 'id': 'call_q9MgGFjqJbV2xSUX93WqxmOt', 'index': 0}])], tool_call_id='call_q9MgGFjqJbV2xSUX93WqxmOt')], 'messages': [AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_q9MgGFjqJbV2xSUX93WqxmOt', 'function': {'arguments': '{\"input\":3}', 'name': 'magic_function'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-c68fd76f-a3c3-4c3c-bfd7-748c171ed4b8', tool_calls=[{'name': 'magic_function', 'args': {'input': 3}, 'id': 'call_q9MgGFjqJbV2xSUX93WqxmOt'}], tool_call_chunks=[{'name': 'magic_function', 'args': '{\"input\":3}', 'id': 'call_q9MgGFjqJbV2xSUX93WqxmOt', 'index': 0}])]}\n",
      "{'steps': [AgentStep(action=ToolAgentAction(tool='magic_function', tool_input={'input': 3}, log=\"\\nInvoking: `magic_function` with `{'input': 3}`\\n\\n\\n\", message_log=[AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_q9MgGFjqJbV2xSUX93WqxmOt', 'function': {'arguments': '{\"input\":3}', 'name': 'magic_function'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-c68fd76f-a3c3-4c3c-bfd7-748c171ed4b8', tool_calls=[{'name': 'magic_function', 'args': {'input': 3}, 'id': 'call_q9MgGFjqJbV2xSUX93WqxmOt'}], tool_call_chunks=[{'name': 'magic_function', 'args': '{\"input\":3}', 'id': 'call_q9MgGFjqJbV2xSUX93WqxmOt', 'index': 0}])], tool_call_id='call_q9MgGFjqJbV2xSUX93WqxmOt'), observation=5)], 'messages': [FunctionMessage(content='5', name='magic_function')]}\n",
      "{'output': 'The value of `magic_function(3)` is 5.', 'messages': [AIMessage(content='The value of `magic_function(3)` is 5.')]}\n"
     ]
    }
   ],
   "source": [
    "from langchain.agents import AgentExecutor, create_tool_calling_agent\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o\")\n",
    "\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant.\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "        # Placeholders fill up a **list** of messages\n",
    "        (\"placeholder\", \"{agent_scratchpad}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "@tool\n",
    "def magic_function(input: int) -> int:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    return input + 2\n",
    "\n",
    "\n",
    "tools = [magic_function]\n",
    "\n",
    "agent = create_tool_calling_agent(model, tools, prompt=prompt)\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools)\n",
    "\n",
    "for step in agent_executor.stream({\"input\": query}):\n",
    "    print(step)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46ccbcbf",
   "metadata": {},
   "source": [
    "#### In LangGraph\n",
    "\n",
    "In LangGraph, things are handled natively using [stream](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.graph.CompiledGraph.stream) or the asynchronous `astream` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "86abbe07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_yTjXXibj76tyFyPRa1soLo0S', 'function': {'arguments': '{\"input\":3}', 'name': 'magic_function'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 70, 'total_tokens': 84}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b275f314-c42e-4e77-9dec-5c23f7dbd53b-0', tool_calls=[{'name': 'magic_function', 'args': {'input': 3}, 'id': 'call_yTjXXibj76tyFyPRa1soLo0S'}])]}}\n",
      "{'tools': {'messages': [ToolMessage(content='5', name='magic_function', id='41c5f227-528d-4483-a313-b03b23b1d327', tool_call_id='call_yTjXXibj76tyFyPRa1soLo0S')]}}\n",
      "{'agent': {'messages': [AIMessage(content='The value of `magic_function(3)` is 5.', response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 93, 'total_tokens': 107}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'stop', 'logprobs': None}, id='run-0ef12b6e-415d-4758-9b62-5e5e1b350072-0')]}}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import AnyMessage\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant.\"),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "def _modify_messages(messages: list[AnyMessage]):\n",
    "    return prompt.invoke({\"messages\": messages}).to_messages()\n",
    "\n",
    "\n",
    "app = create_react_agent(model, tools, messages_modifier=_modify_messages)\n",
    "\n",
    "\n",
    "for step in app.stream({\"messages\": [(\"human\", query)]}, stream_mode=\"updates\"):\n",
    "    print(step)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6898ccbc-42b1-4373-954a-2c7b3849fbb0",
   "metadata": {},
   "source": [
    "## `return_intermediate_steps`\n",
    "\n",
    "Setting this parameter on AgentExecutor allows users to access intermediate_steps, which pairs agent actions (e.g., tool invocations) with their outcomes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4eff44bc-a620-4c8a-97b1-268692a842bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(ToolAgentAction(tool='magic_function', tool_input={'input': 3}, log=\"\\nInvoking: `magic_function` with `{'input': 3}`\\n\\n\\n\", message_log=[AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_ABI4hftfEdnVgKyfF6OzZbca', 'function': {'arguments': '{\"input\":3}', 'name': 'magic_function'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-837e794f-cfd8-40e0-8abc-4d98ced11b75', tool_calls=[{'name': 'magic_function', 'args': {'input': 3}, 'id': 'call_ABI4hftfEdnVgKyfF6OzZbca'}], tool_call_chunks=[{'name': 'magic_function', 'args': '{\"input\":3}', 'id': 'call_ABI4hftfEdnVgKyfF6OzZbca', 'index': 0}])], tool_call_id='call_ABI4hftfEdnVgKyfF6OzZbca'), 5)]\n"
     ]
    }
   ],
   "source": [
    "agent_executor = AgentExecutor(agent=agent, tools=tools, return_intermediate_steps=True)\n",
    "result = agent_executor.invoke({\"input\": query})\n",
    "print(result[\"intermediate_steps\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "594f7567-302f-4fa8-85bb-025ac8322162",
   "metadata": {},
   "source": [
    "By default the [react agent executor](https://langchain-ai.github.io/langgraph/reference/prebuilt/#create_react_agent) in LangGraph appends all messages to the central state. Therefore, it is easy to see any intermediate steps by just looking at the full state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4f4364ea-dffe-4d25-bdce-ef7d0020b880",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='what is the value of magic_function(3)?', id='0f63e437-c4d8-4da9-b6f5-b293ebfe4a64'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_S96v28LlI6hNkQrNnIio0JPh', 'function': {'arguments': '{\"input\":3}', 'name': 'magic_function'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 64, 'total_tokens': 78}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-ffef7898-14b1-4537-ad90-7c000a8a5d25-0', tool_calls=[{'name': 'magic_function', 'args': {'input': 3}, 'id': 'call_S96v28LlI6hNkQrNnIio0JPh'}]),\n",
       "  ToolMessage(content='5', name='magic_function', id='fbd9df4e-1dda-4d3e-9044-b001f7875476', tool_call_id='call_S96v28LlI6hNkQrNnIio0JPh'),\n",
       "  AIMessage(content='The value of `magic_function(3)` is 5.', response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 87, 'total_tokens': 101}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'stop', 'logprobs': None}, id='run-e5d94c54-d9f4-45cd-be8e-a9101a8d88d6-0')]}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "app = create_react_agent(model, tools=tools)\n",
    "\n",
    "messages = app.invoke({\"messages\": [(\"human\", query)]})\n",
    "\n",
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45b528e5-57e1-450e-8d91-513eab53b543",
   "metadata": {},
   "source": [
    "## `max_iterations`\n",
    "\n",
    "`AgentExecutor` implements a `max_iterations` parameter, whereas this is controlled via `recursion_limit` in LangGraph.\n",
    "\n",
    "Note that in AgentExecutor, an \"iteration\" includes a full turn of tool invocation and execution. In LangGraph, each step contributes to the recursion limit, so we will need to multiply by two (and add one) to get equivalent results.\n",
    "\n",
    "If the recursion limit is reached, LangGraph raises a specific exception type, that we can catch and manage similarly to AgentExecutor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "16f189a7-fc78-4cb5-aa16-a94ca06401a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "def magic_function(input: str) -> str:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    return \"Sorry, there was an error. Please try again.\"\n",
    "\n",
    "\n",
    "tools = [magic_function]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c96aefd7-6f6e-4670-aca6-1ac3d4e7871f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `magic_function` with `{'input': '3'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mSorry, there was an error. Please try again.\u001b[0m\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `magic_function` with `{'input': '3'}`\n",
      "responded: Parece que hubo un error al intentar obtener el valor de `magic_function(3)`. Permíteme intentarlo de nuevo.\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mSorry, there was an error. Please try again.\u001b[0m\u001b[32;1m\u001b[1;3mAún no puedo obtener el valor de `magic_function(3)`. ¿Hay algo más en lo que pueda ayudarte?\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'what is the value of magic_function(3)?',\n",
       " 'output': 'Aún no puedo obtener el valor de `magic_function(3)`. ¿Hay algo más en lo que pueda ayudarte?'}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant. Respond only in Spanish.\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "        # Placeholders fill up a **list** of messages\n",
    "        (\"placeholder\", \"{agent_scratchpad}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "agent = create_tool_calling_agent(model, tools, prompt)\n",
    "agent_executor = AgentExecutor(\n",
    "    agent=agent,\n",
    "    tools=tools,\n",
    "    verbose=True,\n",
    "    max_iterations=3,\n",
    ")\n",
    "\n",
    "agent_executor.invoke({\"input\": query})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b974a91f-6ae8-4644-83d9-73666258a6db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('human', 'what is the value of magic_function(3)?')\n",
      "content='' additional_kwargs={'tool_calls': [{'id': 'call_pFdKcCu5taDTtOOfX14vEDRp', 'function': {'arguments': '{\"input\":\"3\"}', 'name': 'magic_function'}, 'type': 'function'}]} response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 64, 'total_tokens': 78}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'tool_calls', 'logprobs': None} id='run-25836468-ba7e-43be-a7cf-76bba06a2a08-0' tool_calls=[{'name': 'magic_function', 'args': {'input': '3'}, 'id': 'call_pFdKcCu5taDTtOOfX14vEDRp'}]\n",
      "content='Sorry, there was an error. Please try again.' name='magic_function' id='1a08b883-9c7b-4969-9e9b-67ce64cdcb5f' tool_call_id='call_pFdKcCu5taDTtOOfX14vEDRp'\n",
      "content='It seems there was an error when trying to apply the magic function. Let me try again.' additional_kwargs={'tool_calls': [{'id': 'call_DA0lpDIkBFg2GHy4WsEcZG4K', 'function': {'arguments': '{\"input\":\"3\"}', 'name': 'magic_function'}, 'type': 'function'}]} response_metadata={'token_usage': {'completion_tokens': 34, 'prompt_tokens': 97, 'total_tokens': 131}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'tool_calls', 'logprobs': None} id='run-d571b774-0ea3-4e35-8b7d-f32932c3f3cc-0' tool_calls=[{'name': 'magic_function', 'args': {'input': '3'}, 'id': 'call_DA0lpDIkBFg2GHy4WsEcZG4K'}]\n",
      "content='Sorry, there was an error. Please try again.' name='magic_function' id='0b45787b-c82a-487f-9a5a-de129c30460f' tool_call_id='call_DA0lpDIkBFg2GHy4WsEcZG4K'\n",
      "content='It appears that there is a consistent issue when trying to apply the magic function to the input \"3.\" This could be due to various reasons, such as the input not being in the correct format or an internal error.\\n\\nIf you have any other questions or if there\\'s something else you\\'d like to try, please let me know!' response_metadata={'token_usage': {'completion_tokens': 66, 'prompt_tokens': 153, 'total_tokens': 219}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'stop', 'logprobs': None} id='run-50a962e6-21b7-4327-8dea-8e2304062627-0'\n"
     ]
    }
   ],
   "source": [
    "from langgraph.errors import GraphRecursionError\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "RECURSION_LIMIT = 2 * 3 + 1\n",
    "\n",
    "app = create_react_agent(model, tools=tools)\n",
    "\n",
    "try:\n",
    "    for chunk in app.stream(\n",
    "        {\"messages\": [(\"human\", query)]},\n",
    "        {\"recursion_limit\": RECURSION_LIMIT},\n",
    "        stream_mode=\"values\",\n",
    "    ):\n",
    "        print(chunk[\"messages\"][-1])\n",
    "except GraphRecursionError:\n",
    "    print({\"input\": query, \"output\": \"Agent stopped due to max iterations.\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a527158-ada5-4774-a98b-8272c6b6b2c0",
   "metadata": {},
   "source": [
    "## `max_execution_time`\n",
    "\n",
    "`AgentExecutor` implements a `max_execution_time` parameter, allowing users to abort a run that exceeds a total time limit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4b8498fc-a7af-4164-a401-d8714f082306",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `magic_function` with `{'input': '3'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mSorry, there was an error. Please try again.\u001b[0m\u001b[32;1m\u001b[1;3m\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'what is the value of magic_function(3)?',\n",
       " 'output': 'Agent stopped due to max iterations.'}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "\n",
    "@tool\n",
    "def magic_function(input: str) -> str:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    time.sleep(2.5)\n",
    "    return \"Sorry, there was an error. Please try again.\"\n",
    "\n",
    "\n",
    "tools = [magic_function]\n",
    "\n",
    "agent = create_tool_calling_agent(model, tools, prompt)\n",
    "agent_executor = AgentExecutor(\n",
    "    agent=agent,\n",
    "    tools=tools,\n",
    "    max_execution_time=2,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "agent_executor.invoke({\"input\": query})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d02eb025",
   "metadata": {},
   "source": [
    "With LangGraph's react agent, you can control timeouts on two levels. \n",
    "\n",
    "You can set a `step_timeout` to bound each **step**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a2b29113-e6be-4f91-aa4c-5c63dea3e423",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_HaQkeCwD5QskzJzFixCBacZ4', 'function': {'arguments': '{\"input\":\"3\"}', 'name': 'magic_function'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 64, 'total_tokens': 78}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-596c9200-771f-436d-8576-72fcb81620f1-0', tool_calls=[{'name': 'magic_function', 'args': {'input': '3'}, 'id': 'call_HaQkeCwD5QskzJzFixCBacZ4'}])]}}\n",
      "------\n",
      "{'input': 'what is the value of magic_function(3)?', 'output': 'Agent stopped due to max iterations.'}\n"
     ]
    }
   ],
   "source": [
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "app = create_react_agent(model, tools=tools)\n",
    "# Set the max timeout for each step here\n",
    "app.step_timeout = 2\n",
    "\n",
    "try:\n",
    "    for chunk in app.stream({\"messages\": [(\"human\", query)]}):\n",
    "        print(chunk)\n",
    "        print(\"------\")\n",
    "except TimeoutError:\n",
    "    print({\"input\": query, \"output\": \"Agent stopped due to max iterations.\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32a9db70",
   "metadata": {},
   "source": [
    "The other way to set a single max timeout for an entire run is to directly use the python stdlib [asyncio](https://docs.python.org/3/library/asyncio.html) library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e9eb55f4-a321-4bac-b52d-9e43b411cf92",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_4agJXUHtmHrOOMogjF6ZuzAv', 'function': {'arguments': '{\"input\":\"3\"}', 'name': 'magic_function'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 64, 'total_tokens': 78}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-a1c77db7-405f-43d9-8d57-751f2ca1a58c-0', tool_calls=[{'name': 'magic_function', 'args': {'input': '3'}, 'id': 'call_4agJXUHtmHrOOMogjF6ZuzAv'}])]}}\n",
      "------\n",
      "Task Cancelled.\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "app = create_react_agent(model, tools=tools)\n",
    "\n",
    "\n",
    "async def stream(app, inputs):\n",
    "    async for chunk in app.astream({\"messages\": [(\"human\", query)]}):\n",
    "        print(chunk)\n",
    "        print(\"------\")\n",
    "\n",
    "\n",
    "try:\n",
    "    task = asyncio.create_task(stream(app, {\"messages\": [(\"human\", query)]}))\n",
    "    await asyncio.wait_for(task, timeout=3)\n",
    "except TimeoutError:\n",
    "    print(\"Task Cancelled.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4884ac87",
   "metadata": {},
   "source": [
    "## `early_stopping_method`\n",
    "\n",
    "With LangChain's [AgentExecutor](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor.iter), you could configure an [early_stopping_method](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor.early_stopping_method) to either return a string saying \"Agent stopped due to iteration limit or time limit.\" (`\"force\"`) or prompt the LLM a final time to respond (`\"generate\"`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3f6e2cf2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output with early_stopping_method='force':\n",
      "Agent stopped due to max iterations.\n"
     ]
    }
   ],
   "source": [
    "from langchain.agents import AgentExecutor, create_tool_calling_agent\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o\")\n",
    "\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant.\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "        # Placeholders fill up a **list** of messages\n",
    "        (\"placeholder\", \"{agent_scratchpad}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "@tool\n",
    "def magic_function(input: int) -> int:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    return \"Sorry there was an error, please try again.\"\n",
    "\n",
    "\n",
    "tools = [magic_function]\n",
    "\n",
    "agent = create_tool_calling_agent(model, tools, prompt=prompt)\n",
    "agent_executor = AgentExecutor(\n",
    "    agent=agent, tools=tools, early_stopping_method=\"force\", max_iterations=1\n",
    ")\n",
    "\n",
    "result = agent_executor.invoke({\"input\": query})\n",
    "print(\"Output with early_stopping_method='force':\")\n",
    "print(result[\"output\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "706e05c4",
   "metadata": {},
   "source": [
    "#### In LangGraph\n",
    "\n",
    "In LangGraph, you can explicitly handle the response behavior outside the agent, since the full state can be accessed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "73cabbc4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('human', 'what is the value of magic_function(3)?')\n",
      "content='' additional_kwargs={'tool_calls': [{'id': 'call_bTURmOn9C8zslmn0kMFeykIn', 'function': {'arguments': '{\"input\":3}', 'name': 'magic_function'}, 'type': 'function'}]} response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 64, 'total_tokens': 78}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'tool_calls', 'logprobs': None} id='run-0844a504-7e6b-4ea6-a069-7017e38121ee-0' tool_calls=[{'name': 'magic_function', 'args': {'input': 3}, 'id': 'call_bTURmOn9C8zslmn0kMFeykIn'}]\n",
      "content='Sorry there was an error, please try again.' name='magic_function' id='00d5386f-eb23-4628-9a29-d9ce6a7098cc' tool_call_id='call_bTURmOn9C8zslmn0kMFeykIn'\n",
      "content='' additional_kwargs={'tool_calls': [{'id': 'call_JYqvvvWmXow2u012DuPoDHFV', 'function': {'arguments': '{\"input\":3}', 'name': 'magic_function'}, 'type': 'function'}]} response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 96, 'total_tokens': 110}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_729ea513f7', 'finish_reason': 'tool_calls', 'logprobs': None} id='run-b73b1b1c-c829-4348-98cd-60b315c85448-0' tool_calls=[{'name': 'magic_function', 'args': {'input': 3}, 'id': 'call_JYqvvvWmXow2u012DuPoDHFV'}]\n",
      "{'input': 'what is the value of magic_function(3)?', 'output': 'Agent stopped due to max iterations.'}\n"
     ]
    }
   ],
   "source": [
    "from langgraph.errors import GraphRecursionError\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "RECURSION_LIMIT = 2 * 1 + 1\n",
    "\n",
    "app = create_react_agent(model, tools=tools)\n",
    "\n",
    "try:\n",
    "    for chunk in app.stream(\n",
    "        {\"messages\": [(\"human\", query)]},\n",
    "        {\"recursion_limit\": RECURSION_LIMIT},\n",
    "        stream_mode=\"values\",\n",
    "    ):\n",
    "        print(chunk[\"messages\"][-1])\n",
    "except GraphRecursionError:\n",
    "    print({\"input\": query, \"output\": \"Agent stopped due to max iterations.\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "017fe20e",
   "metadata": {},
   "source": [
    "## `trim_intermediate_steps`\n",
    "\n",
    "With LangChain's [AgentExecutor](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor), you could trim the intermediate steps of long-running agents using [trim_intermediate_steps](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html#langchain.agents.agent.AgentExecutor.trim_intermediate_steps), which is either an integer (indicating the agent should keep the last N steps) or a custom function.\n",
    "\n",
    "For instance, we could trim the value so the agent only sees the most recent intermediate step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b94bb169",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Call number: 1\n",
      "Call number: 2\n",
      "Call number: 3\n",
      "Call number: 4\n",
      "Call number: 5\n",
      "Call number: 6\n",
      "Call number: 7\n",
      "Call number: 8\n",
      "Call number: 9\n",
      "Call number: 10\n",
      "Call number: 11\n",
      "Call number: 12\n",
      "Call number: 13\n",
      "Call number: 14\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stopping agent prematurely due to triggering stop condition\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Call number: 15\n"
     ]
    }
   ],
   "source": [
    "from langchain.agents import AgentExecutor, create_tool_calling_agent\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o\")\n",
    "\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant.\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "        # Placeholders fill up a **list** of messages\n",
    "        (\"placeholder\", \"{agent_scratchpad}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "magic_step_num = 1\n",
    "\n",
    "\n",
    "@tool\n",
    "def magic_function(input: int) -> int:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    global magic_step_num\n",
    "    print(f\"Call number: {magic_step_num}\")\n",
    "    magic_step_num += 1\n",
    "    return input + magic_step_num\n",
    "\n",
    "\n",
    "tools = [magic_function]\n",
    "\n",
    "agent = create_tool_calling_agent(model, tools, prompt=prompt)\n",
    "\n",
    "\n",
    "def trim_steps(steps: list):\n",
    "    # Let's give the agent amnesia\n",
    "    return []\n",
    "\n",
    "\n",
    "agent_executor = AgentExecutor(\n",
    "    agent=agent, tools=tools, trim_intermediate_steps=trim_steps\n",
    ")\n",
    "\n",
    "\n",
    "query = \"Call the magic function 4 times in sequence with the value 3. You cannot call it multiple times at once.\"\n",
    "\n",
    "for step in agent_executor.stream({\"input\": query}):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d450c5a",
   "metadata": {},
   "source": [
    "#### In LangGraph\n",
    "\n",
    "We can use the [`messages_modifier`](https://langchain-ai.github.io/langgraph/reference/prebuilt/#create_react_agent) just as before when passing in [prompt templates](#prompt-templates)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b309ba9a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Call number: 1\n",
      "Call number: 2\n",
      "Call number: 3\n",
      "Call number: 4\n",
      "Call number: 5\n",
      "Call number: 6\n",
      "Call number: 7\n",
      "Call number: 8\n",
      "Call number: 9\n",
      "Call number: 10\n",
      "Call number: 11\n",
      "Call number: 12\n",
      "Stopping agent prematurely due to triggering stop condition\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import AnyMessage\n",
    "from langgraph.errors import GraphRecursionError\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "magic_step_num = 1\n",
    "\n",
    "\n",
    "@tool\n",
    "def magic_function(input: int) -> int:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    global magic_step_num\n",
    "    print(f\"Call number: {magic_step_num}\")\n",
    "    magic_step_num += 1\n",
    "    return input + magic_step_num\n",
    "\n",
    "\n",
    "tools = [magic_function]\n",
    "\n",
    "\n",
    "def _modify_messages(messages: list[AnyMessage]):\n",
    "    # Give the agent amnesia, only keeping the original user query\n",
    "    return [(\"system\", \"You are a helpful assistant\"), messages[0]]\n",
    "\n",
    "\n",
    "app = create_react_agent(model, tools, messages_modifier=_modify_messages)\n",
    "\n",
    "try:\n",
    "    for step in app.stream({\"messages\": [(\"human\", query)]}, stream_mode=\"updates\"):\n",
    "        pass\n",
    "except GraphRecursionError as e:\n",
    "    print(\"Stopping agent prematurely due to triggering stop condition\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
