From a391e3f5d32c426d98f8a2a4bc17c731d9664a73 Mon Sep 17 00:00:00 2001 From: Bagatur Date: Mon, 5 Aug 2024 16:51:00 -0700 Subject: [PATCH] fmt --- libs/core/langchain_core/tools.py | 4 +++- .../integration_tests/chat_models.py | 8 +++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/libs/core/langchain_core/tools.py b/libs/core/langchain_core/tools.py index 702cd7b6fe789..53c53690a3776 100644 --- a/libs/core/langchain_core/tools.py +++ b/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, diff --git a/libs/standard-tests/langchain_standard_tests/integration_tests/chat_models.py b/libs/standard-tests/langchain_standard_tests/integration_tests/chat_models.py index 5c23376a91703..b22844fa4d165 100644 --- a/libs/standard-tests/langchain_standard_tests/integration_tests/chat_models.py +++ b/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