diff --git a/docs/docs/how_to/custom_tools.ipynb b/docs/docs/how_to/custom_tools.ipynb
index 0e1262841a1..81853e17a0d 100644
--- a/docs/docs/how_to/custom_tools.ipynb
+++ b/docs/docs/how_to/custom_tools.ipynb
@@ -492,7 +492,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "id": "1dad8f8e",
    "metadata": {
     "execution": {
@@ -504,13 +504,14 @@
    },
    "outputs": [],
    "source": [
-    "from typing import Optional, Type\n",
+    "from typing import Optional\n",
     "\n",
     "from langchain_core.callbacks import (\n",
     "    AsyncCallbackManagerForToolRun,\n",
     "    CallbackManagerForToolRun,\n",
     ")\n",
     "from langchain_core.tools import BaseTool\n",
+    "from langchain_core.tools.base import ArgsSchema\n",
     "from pydantic import BaseModel, Field\n",
     "\n",
     "\n",
@@ -524,7 +525,7 @@
     "class CustomCalculatorTool(BaseTool):\n",
     "    name: str = \"Calculator\"\n",
     "    description: str = \"useful for when you need to answer questions about math\"\n",
-    "    args_schema: Type[BaseModel] = CalculatorInput\n",
+    "    args_schema: Optional[ArgsSchema] = CalculatorInput\n",
     "    return_direct: bool = True\n",
     "\n",
     "    def _run(\n",