diff --git a/docs/modules/agents/tools/custom_tools.ipynb b/docs/modules/agents/tools/custom_tools.ipynb
index 2e2fa6f7..b34cd90d 100644
--- a/docs/modules/agents/tools/custom_tools.ipynb
+++ b/docs/modules/agents/tools/custom_tools.ipynb
@@ -191,6 +191,8 @@
    },
    "outputs": [],
    "source": [
+    "from typing import Type\n",
+    "\n",
     "class CustomSearchTool(BaseTool):\n",
     "    name = \"Search\"\n",
     "    description = \"useful for when you need to answer questions about current events\"\n",
@@ -206,7 +208,7 @@
     "class CustomCalculatorTool(BaseTool):\n",
     "    name = \"Calculator\"\n",
     "    description = \"useful for when you need to answer questions about math\"\n",
-    "    args_schema=CalculatorInput\n",
+    "    args_schema: Type[BaseModel] = CalculatorInput\n",
     "\n",
     "    def _run(self, query: str) -> str:\n",
     "        \"\"\"Use the tool.\"\"\"\n",