test_args_schema_explicitly_typed() — langchain Function Reference

Source Code

libs/core/tests/unit_tests/test_tools.py lines 2001–2043

def test_args_schema_explicitly_typed() -> None:
    """This should test that one can type the args schema as a Pydantic model.

    Please note that this will test using pydantic 2 even though `BaseTool`
    is a Pydantic 1 model!
    """

    class Foo(BaseModel):
        a: int
        b: str

    class SomeTool(BaseTool):
        # type ignoring here since we're allowing overriding a type
        # signature of pydantic.v1.BaseModel with pydantic.BaseModel
        # for pydantic 2!
        args_schema: type[BaseModel] = Foo

        @override
        def _run(self, *args: Any, **kwargs: Any) -> str:
            return "foo"

    tool = SomeTool(name="some_tool", description="some description")

    assert tool.get_input_schema().model_json_schema() == {
        "properties": {
            "a": {"title": "A", "type": "integer"},
            "b": {"title": "B", "type": "string"},
        },
        "required": ["a", "b"],
        "title": "Foo",
        "type": "object",
    }

    assert _get_tool_call_json_schema(tool) == {
        "description": "some description",
        "properties": {
            "a": {"title": "A", "type": "integer"},
            "b": {"title": "B", "type": "string"},
        },
        "required": ["a", "b"],
        "title": "some_tool",
        "type": "object",
    }