Cannot access ToolRuntime in BaseTool subclass

Adding a ToolRuntime to the overriden _run method’s signature does help to inject ToolRuntime, which reports error as

TypeError: Multipler._run() missing 1 required positional argument: 'runtime'

Codes reproducing the error goes as

from typing import Type

from langchain.tools import ToolRuntime
from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field


class MultiplyInput(BaseModel):
    a: int = Field(description="first number")
    b: int = Field(description="second number")


class Multipler(BaseTool):
    name: str = "Multiplier"
    description: str = "a tool for multiplying two numbers"
    args_schema: Type[BaseModel] = MultiplyInput
    return_direct: bool = False

    def _run(self, a: int, b: int, runtime: ToolRuntime) -> int:
        print(f"runtime = {runtime}")

        return a * b


if __name__ == "__main__":
    import os

    import dotenv
    from langchain.agents import create_agent
    from langchain_openai import ChatOpenAI

    # load OPENAI_API_KEY, OPENAI_API_BASE_URL and OPENAI_MODEL from .env
    dotenv.load_dotenv()

    model = ChatOpenAI(
        api_key=os.environ["OPENAI_API_KEY"],
        base_url=os.environ["OPENAI_API_BASE_URL"],
        model=os.environ["OPENAI_MODEL"],
    )

    agent = create_agent(
        model, tools=[Multipler()], system_prompt="You are a helpful assistant."
    )

    r = agent.invoke({"messages": [{"role": "user", "content": "calculate 2 * 3"}]})

    for v in r["messages"]:
        v.pretty_print()

Really appreciate if someone can help me out

hi @sammyne

when you use args_schema and want to access the runtime, all you need is to define it on the args schema class like this:

class MultiplyInput(BaseModel):
    a: int = Field(description="first number")
    b: int = Field(description="second number")
    runtime: ToolRuntime  # injected by ToolNode, internally

which means the class must be aligned with the arguments of the _run method.

See that post How to access tool call id using runtime - #2 by pawel-twardziak

1 Like