args_schema type hint on subclassing (#3323)

per #3297 Co-authored-by: Daniel Chalef <daniel.chalef@private.org>
langchain-ai · Apr 26, 2023 · 6a0abcc · 6a0abcc
1 parent beb0f6f
commit 6a0abcc
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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",