test_tool_allows_extra_runtime_args_with_custom_schema() — langchain Function Reference

Source Code

libs/core/tests/unit_tests/test_tools.py lines 2414–2441

def test_tool_allows_extra_runtime_args_with_custom_schema(
    schema_format: Literal["model", "json_schema"],
) -> None:
    """Ensure runtime args are preserved even if not in the args schema."""

    class InputSchema(BaseModel):
        query: str

    captured: dict[str, Any] = {}

    @dataclass
    class MyRuntime(_DirectlyInjectedToolArg):
        some_obj: object

    args_schema = (
        InputSchema if schema_format == "model" else InputSchema.model_json_schema()
    )

    @tool(args_schema=args_schema)
    def runtime_tool(query: str, runtime: MyRuntime) -> str:
        """Echo the query and capture runtime value."""
        captured["runtime"] = runtime
        return query

    runtime_obj = object()
    runtime = MyRuntime(some_obj=runtime_obj)
    assert runtime_tool.invoke({"query": "hello", "runtime": runtime}) == "hello"
    assert captured["runtime"] is runtime