{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3c284df8",
   "metadata": {},
   "source": [
    "# XML Agent\n",
    "\n",
    "Some language models (like Anthropic's Claude) are particularly good at reasoning/writing XML. This goes over how to use an agent that uses XML when prompting. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe972808",
   "metadata": {},
   "source": [
    "## Initialize the tools\n",
    "\n",
    "We will initialize some fake tools for demo purposes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ba547497",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import tool\n",
    "\n",
    "\n",
    "@tool\n",
    "def search(query: str) -> str:\n",
    "    \"\"\"Search things about current events.\"\"\"\n",
    "    return \"32 degrees\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e30e99e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [search]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "401db6ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatAnthropic\n",
    "\n",
    "model = ChatAnthropic(model=\"claude-2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90f83099",
   "metadata": {},
   "source": [
    "## Use LangChain Expression Language\n",
    "\n",
    "We will first show how to create this agent using LangChain Expression Language"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "78937679",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import hub\n",
    "from langchain.agents.format_scratchpad import format_xml\n",
    "from langchain.agents.output_parsers import XMLAgentOutputParser\n",
    "from langchain.tools.render import render_text_description"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "54fc5a22",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = hub.pull(\"hwchase17/xml-agent\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b1802fcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = prompt.partial(\n",
    "    tools=render_text_description(tools),\n",
    "    tool_names=\", \".join([t.name for t in tools]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f9d2ead2",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_with_stop = model.bind(stop=[\"</tool_input>\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ebadf04f",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = (\n",
    "    {\n",
    "        \"question\": lambda x: x[\"question\"],\n",
    "        \"agent_scratchpad\": lambda x: format_xml(x[\"intermediate_steps\"]),\n",
    "    }\n",
    "    | prompt\n",
    "    | llm_with_stop\n",
    "    | XMLAgentOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e2bb03e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6ce9f9a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e14affef",
   "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 <tool>search</tool>\n",
      "<tool_input>weather in new york\u001b[0m\u001b[36;1m\u001b[1;3m32 degrees\u001b[0m\u001b[32;1m\u001b[1;3m <tool>search</tool>\n",
      "<tool_input>weather in new york\u001b[0m\u001b[36;1m\u001b[1;3m32 degrees\u001b[0m\u001b[32;1m\u001b[1;3m <final_answer>\n",
      "The weather in New York is 32 degrees.\n",
      "</final_answer>\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'question': 'whats the weather in New york?',\n",
       " 'output': '\\nThe weather in New York is 32 degrees.\\n'}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.invoke({\"question\": \"whats the weather in New york?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42ff473d",
   "metadata": {},
   "source": [
    "## Use off-the-shelf agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "7e5e73e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import XMLAgent\n",
    "from langchain.chains import LLMChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "2d8454be",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = LLMChain(\n",
    "    llm=model,\n",
    "    prompt=XMLAgent.get_default_prompt(),\n",
    "    output_parser=XMLAgent.get_default_output_parser(),\n",
    ")\n",
    "agent = XMLAgent(tools=tools, llm_chain=chain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "bca6096f",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "71b872b1",
   "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 <tool>search</tool>\n",
      "<tool_input>weather in new york\u001b[0m\u001b[36;1m\u001b[1;3m32 degrees\u001b[0m\u001b[32;1m\u001b[1;3m\n",
      "\n",
      "<final_answer>The weather in New York is 32 degrees\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'whats the weather in New york?',\n",
       " 'output': 'The weather in New York is 32 degrees'}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.invoke({\"input\": \"whats the weather in New york?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cca87246",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
