fmt

libs/core/langchain_core/tools.py

@@ -1055,10 +1055,12 @@ def add(a: int, b: int) -> int:
         )
 
 
+# TODO: Type args_schema as TypeBaseModel if we can get mypy to correctly recognize
+# pydantic v2 BaseModel classes.
 def tool(
     *args: Union[str, Callable, Runnable],
     return_direct: bool = False,
-    args_schema: Optional[TypeBaseModel] = None,
+    args_schema: Optional[Type] = None,
     infer_schema: bool = True,
     response_format: Literal["content", "content_and_artifact"] = "content",
     parse_docstring: bool = False,

libs/standard-tests/langchain_standard_tests/integration_tests/chat_models.py

@@ -18,6 +18,8 @@
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain_core.tools import tool
+from pydantic import BaseModel as RawBaseModel
+from pydantic import Field as RawField
 
 from langchain_standard_tests.unit_tests.chat_models import (
     ChatModelTests,
@@ -26,7 +28,11 @@
 from langchain_standard_tests.utils.pydantic import PYDANTIC_MAJOR_VERSION
 
 
-@tool
+class MagicFunctionSchema(RawBaseModel):
+    input: int = RawField(..., gt=-1000, lt=1000)
+
+
+@tool(args_schema=MagicFunctionSchema)
 def magic_function(input: int) -> int:
     """Applies a magic function to an input."""
     return input + 2