docs/modules/agents/tools/tool_input_validation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Tool Input Schema\n",
    "\n",
    "By default, tools infer the argument schema by inspecting the function signature. For more strict requirements, custom input schema can be specified, along with custom validation logic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from typing import Any, Dict\n",
    "\n",
    "from langchain.agents import AgentType, initialize_agent\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.tools.requests.tool import RequestsGetTool, TextRequestsWrapper\n",
    "from pydantic import BaseModel, Field, root_validator\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install tldextract > /dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import tldextract\n",
    "\n",
    "_APPROVED_DOMAINS = {\n",
    "    \"langchain\",\n",
    "    \"wikipedia\",\n",
    "}\n",
    "\n",
    "class ToolInputSchema(BaseModel):\n",
    "\n",
    "    url: str = Field(...)\n",
    "    \n",
    "    @root_validator\n",
    "    def validate_query(cls, values: Dict[str, Any]) -> Dict:\n",
    "        url = values[\"url\"]\n",
    "        domain = tldextract.extract(url).domain\n",
    "        if domain not in _APPROVED_DOMAINS:\n",
    "            raise ValueError(f\"Domain {domain} is not on the approved list:\"\n",
    "                             f\" {sorted(_APPROVED_DOMAINS)}\")\n",
    "        return values\n",
    "    \n",
    "tool = RequestsGetTool(args_schema=ToolInputSchema, requests_wrapper=TextRequestsWrapper())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent = initialize_agent([tool], llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The main title of langchain.com is \"LANG CHAIN 🦜️🔗 Official Home Page\"\n"
     ]
    }
   ],
   "source": [
    "# This will succeed, since there aren't any arguments that will be triggered during validation\n",
    "answer = agent.run(\"What's the main title on langchain.com?\")\n",
    "print(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "ename": "ValidationError",
     "evalue": "1 validation error for ToolInputSchema\n__root__\n  Domain google is not on the approved list: ['langchain', 'wikipedia'] (type=value_error)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValidationError\u001b[0m                           Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[7], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m agent\u001b[39m.\u001b[39;49mrun(\u001b[39m\"\u001b[39;49m\u001b[39mWhat\u001b[39;49m\u001b[39m'\u001b[39;49m\u001b[39ms the main title on google.com?\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n",
      "File \u001b[0;32m~/code/lc/lckg/langchain/chains/base.py:213\u001b[0m, in \u001b[0;36mChain.run\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    211\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mlen\u001b[39m(args) \u001b[39m!=\u001b[39m \u001b[39m1\u001b[39m:\n\u001b[1;32m    212\u001b[0m         \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\u001b[39m\"\u001b[39m\u001b[39m`run` supports only one positional argument.\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m--> 213\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m(args[\u001b[39m0\u001b[39;49m])[\u001b[39mself\u001b[39m\u001b[39m.\u001b[39moutput_keys[\u001b[39m0\u001b[39m]]\n\u001b[1;32m    215\u001b[0m \u001b[39mif\u001b[39;00m kwargs \u001b[39mand\u001b[39;00m \u001b[39mnot\u001b[39;00m args:\n\u001b[1;32m    216\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m(kwargs)[\u001b[39mself\u001b[39m\u001b[39m.\u001b[39moutput_keys[\u001b[39m0\u001b[39m]]\n",
      "File \u001b[0;32m~/code/lc/lckg/langchain/chains/base.py:116\u001b[0m, in \u001b[0;36mChain.__call__\u001b[0;34m(self, inputs, return_only_outputs)\u001b[0m\n\u001b[1;32m    114\u001b[0m \u001b[39mexcept\u001b[39;00m (\u001b[39mKeyboardInterrupt\u001b[39;00m, \u001b[39mException\u001b[39;00m) \u001b[39mas\u001b[39;00m e:\n\u001b[1;32m    115\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcallback_manager\u001b[39m.\u001b[39mon_chain_error(e, verbose\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mverbose)\n\u001b[0;32m--> 116\u001b[0m     \u001b[39mraise\u001b[39;00m e\n\u001b[1;32m    117\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcallback_manager\u001b[39m.\u001b[39mon_chain_end(outputs, verbose\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mverbose)\n\u001b[1;32m    118\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprep_outputs(inputs, outputs, return_only_outputs)\n",
      "File \u001b[0;32m~/code/lc/lckg/langchain/chains/base.py:113\u001b[0m, in \u001b[0;36mChain.__call__\u001b[0;34m(self, inputs, return_only_outputs)\u001b[0m\n\u001b[1;32m    107\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcallback_manager\u001b[39m.\u001b[39mon_chain_start(\n\u001b[1;32m    108\u001b[0m     {\u001b[39m\"\u001b[39m\u001b[39mname\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__class__\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m},\n\u001b[1;32m    109\u001b[0m     inputs,\n\u001b[1;32m    110\u001b[0m     verbose\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mverbose,\n\u001b[1;32m    111\u001b[0m )\n\u001b[1;32m    112\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 113\u001b[0m     outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call(inputs)\n\u001b[1;32m    114\u001b[0m \u001b[39mexcept\u001b[39;00m (\u001b[39mKeyboardInterrupt\u001b[39;00m, \u001b[39mException\u001b[39;00m) \u001b[39mas\u001b[39;00m e:\n\u001b[1;32m    115\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcallback_manager\u001b[39m.\u001b[39mon_chain_error(e, verbose\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mverbose)\n",
      "File \u001b[0;32m~/code/lc/lckg/langchain/agents/agent.py:792\u001b[0m, in \u001b[0;36mAgentExecutor._call\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m    790\u001b[0m \u001b[39m# We now enter the agent loop (until it returns something).\u001b[39;00m\n\u001b[1;32m    791\u001b[0m \u001b[39mwhile\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_should_continue(iterations, time_elapsed):\n\u001b[0;32m--> 792\u001b[0m     next_step_output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_take_next_step(\n\u001b[1;32m    793\u001b[0m         name_to_tool_map, color_mapping, inputs, intermediate_steps\n\u001b[1;32m    794\u001b[0m     )\n\u001b[1;32m    795\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39misinstance\u001b[39m(next_step_output, AgentFinish):\n\u001b[1;32m    796\u001b[0m         \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_return(next_step_output, intermediate_steps)\n",
      "File \u001b[0;32m~/code/lc/lckg/langchain/agents/agent.py:695\u001b[0m, in \u001b[0;36mAgentExecutor._take_next_step\u001b[0;34m(self, name_to_tool_map, color_mapping, inputs, intermediate_steps)\u001b[0m\n\u001b[1;32m    693\u001b[0m         tool_run_kwargs[\u001b[39m\"\u001b[39m\u001b[39mllm_prefix\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    694\u001b[0m     \u001b[39m# We then call the tool on the tool input to get an observation\u001b[39;00m\n\u001b[0;32m--> 695\u001b[0m     observation \u001b[39m=\u001b[39m tool\u001b[39m.\u001b[39;49mrun(\n\u001b[1;32m    696\u001b[0m         agent_action\u001b[39m.\u001b[39;49mtool_input,\n\u001b[1;32m    697\u001b[0m         verbose\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mverbose,\n\u001b[1;32m    698\u001b[0m         color\u001b[39m=\u001b[39;49mcolor,\n\u001b[1;32m    699\u001b[0m         \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mtool_run_kwargs,\n\u001b[1;32m    700\u001b[0m     )\n\u001b[1;32m    701\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    702\u001b[0m     tool_run_kwargs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39magent\u001b[39m.\u001b[39mtool_run_logging_kwargs()\n",
      "File \u001b[0;32m~/code/lc/lckg/langchain/tools/base.py:110\u001b[0m, in \u001b[0;36mBaseTool.run\u001b[0;34m(self, tool_input, verbose, start_color, color, **kwargs)\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mrun\u001b[39m(\n\u001b[1;32m    102\u001b[0m     \u001b[39mself\u001b[39m,\n\u001b[1;32m    103\u001b[0m     tool_input: Union[\u001b[39mstr\u001b[39m, Dict],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    107\u001b[0m     \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs: Any,\n\u001b[1;32m    108\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mstr\u001b[39m:\n\u001b[1;32m    109\u001b[0m \u001b[39m    \u001b[39m\u001b[39m\"\"\"Run the tool.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 110\u001b[0m     run_input \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_parse_input(tool_input)\n\u001b[1;32m    111\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mverbose \u001b[39mand\u001b[39;00m verbose \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    112\u001b[0m         verbose_ \u001b[39m=\u001b[39m verbose\n",
      "File \u001b[0;32m~/code/lc/lckg/langchain/tools/base.py:71\u001b[0m, in \u001b[0;36mBaseTool._parse_input\u001b[0;34m(self, tool_input)\u001b[0m\n\u001b[1;32m     69\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39missubclass\u001b[39m(input_args, BaseModel):\n\u001b[1;32m     70\u001b[0m     key_ \u001b[39m=\u001b[39m \u001b[39mnext\u001b[39m(\u001b[39miter\u001b[39m(input_args\u001b[39m.\u001b[39m__fields__\u001b[39m.\u001b[39mkeys()))\n\u001b[0;32m---> 71\u001b[0m     input_args\u001b[39m.\u001b[39;49mparse_obj({key_: tool_input})\n\u001b[1;32m     72\u001b[0m \u001b[39m# Passing as a positional argument is more straightforward for\u001b[39;00m\n\u001b[1;32m     73\u001b[0m \u001b[39m# backwards compatability\u001b[39;00m\n\u001b[1;32m     74\u001b[0m \u001b[39mreturn\u001b[39;00m tool_input\n",
      "File \u001b[0;32m~/code/lc/lckg/.venv/lib/python3.11/site-packages/pydantic/main.py:526\u001b[0m, in \u001b[0;36mpydantic.main.BaseModel.parse_obj\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/code/lc/lckg/.venv/lib/python3.11/site-packages/pydantic/main.py:341\u001b[0m, in \u001b[0;36mpydantic.main.BaseModel.__init__\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mValidationError\u001b[0m: 1 validation error for ToolInputSchema\n__root__\n  Domain google is not on the approved list: ['langchain', 'wikipedia'] (type=value_error)"
     ]
    }
   ],
   "source": [
    "agent.run(\"What's the main title on google.com?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}