diff --git a/README.md b/README.md index 45516e68..26cc01ea 100644 --- a/README.md +++ b/README.md @@ -295,6 +295,10 @@ We will define a single tool that prints a story. Tools may optionally take a fi exposed to the agent as a parameter but will be injected by the environment (if part of the function signature). +A tool may also declare a `tool_call_id: str` argument. The environment injects the ID of +the tool call being executed, so a tool can key its work by call. Like `state`, this +argument stays hidden from the agent. + ```py def print_story(story: str, state: ExampleState): """Print a story. diff --git a/src/aviary/env.py b/src/aviary/env.py index 517c27c3..acd9e09f 100644 --- a/src/aviary/env.py +++ b/src/aviary/env.py @@ -238,11 +238,14 @@ async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage: f" { {t.info.name for t in self.tools} }." ) from exc - # we do a special convenience to make - # state be optional in the function signature + # we do a special convenience to make state and tool_call_id be + # optional in the function signature: state is dropped when undeclared, + # tool_call_id is injected when declared. Tool.from_function keeps both + # out of the LLM-facing schema, and our id wins over an LLM-emitted one + fn_parameters = inspect.signature(tool._tool_fn).parameters need_to_filter = ( "state" in function_kwargs - and "state" not in inspect.signature(tool._tool_fn).parameters + and "state" not in fn_parameters and not hasattr(tool._tool_fn, "requires_state") ) filtered_kwargs = ( @@ -250,6 +253,11 @@ async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage: if need_to_filter else function_kwargs ) + tool_call_args = tool_call.function.arguments | ( + {"tool_call_id": tool_call.id} + if "tool_call_id" in fn_parameters + else {} + ) concurrency_context = ( concurrency_lock.read_lock() @@ -262,18 +270,14 @@ async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage: async with concurrency_context: if is_coroutine_callable(tool._tool_fn): content = await maybe_wait_for( - tool._tool_fn( - **tool_call.function.arguments, **filtered_kwargs - ), + tool._tool_fn(**tool_call_args, **filtered_kwargs), exec_timeout, ) else: # If the function is synchronous, run on a thread content = await maybe_wait_for( asyncio.to_thread( - tool._tool_fn, - **tool_call.function.arguments, - **filtered_kwargs, + tool._tool_fn, **tool_call_args, **filtered_kwargs ), exec_timeout, ) diff --git a/src/aviary/tools/argref.py b/src/aviary/tools/argref.py index 4e1ba8ca..756553e9 100644 --- a/src/aviary/tools/argref.py +++ b/src/aviary/tools/argref.py @@ -118,7 +118,7 @@ def argref_by_name( # noqa: PLR0915 >>> # Equivalent to my_func(state.refs["a"], state.refs["b"]) >>> wrapped_fxn("a", "b", state=state) # doctest: +SKIP """ - args_to_skip = (args_to_skip or set()) | {"state", "return"} + args_to_skip = (args_to_skip or set()) | {"state", "tool_call_id", "return"} def decorator(func): # noqa: PLR0915 def get_call_args(*args, **kwargs): diff --git a/src/aviary/tools/base.py b/src/aviary/tools/base.py index 78f74c81..ef5a09fd 100644 --- a/src/aviary/tools/base.py +++ b/src/aviary/tools/base.py @@ -415,8 +415,8 @@ def from_function( required: dict[str, bool] = {} annotations = function.__annotations__ for pname, parameter in inspect.signature(function).parameters.items(): - if pname == "state": - # NOTE: ToolRequestMessage passes state for us, not the LLM + if pname in {"state", "tool_call_id"}: + # NOTE: exec_tool_calls injects state and tool_call_id for us, not the LLM continue d = next( ( diff --git a/tests/test_tools.py b/tests/test_tools.py index 2d18dd2f..4e9bf529 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -19,6 +19,7 @@ from aviary.core import ( INVALID_TOOL_NAME, DummyEnv, + DummyEnvState, Environment, FunctionInfo, Message, @@ -702,6 +703,106 @@ def get_todo_list_no_args(): new_messages = await dummy_env.exec_tool_calls(action) assert new_messages[0].content == "Go for a walk" + @pytest.mark.asyncio + async def test_tool_call_id_injection( + self, dummy_env: DummyEnv, subtests: pytest.Subtests + ) -> None: + # NOTE: tool_call_id is left out of every docstring below, to confirm + # from_function doesn't demand a description for an injected parameter + async def remember(x: int, tool_call_id: str) -> str: # noqa: D417 + """Remember a number. + + Args: + x: Number to remember. + """ + await asyncio.sleep(0.01) # Force concurrent calls to overlap + return f"{tool_call_id}:{x}" + + def remember_sync(x: int, tool_call_id: str) -> str: # noqa: D417 + """Remember a number. + + Args: + x: Number to remember. + """ + return f"{tool_call_id}:{x}" + + def remember_in_state( # noqa: D417 + x: int, state: DummyEnvState, tool_call_id: str + ) -> str: + """Remember a number in the state. + + Args: + x: Number to remember. + """ + state.reward = x + return f"{tool_call_id}:{x}" + + class Rememberer: + async def remember_method(self, x: int, tool_call_id: str) -> str: # noqa: D417 + """Remember a number. + + Args: + x: Number to remember. + """ + return f"{tool_call_id}:{x}" + + tools = { + fn.__name__: Tool.from_function(fn) + for fn in ( + remember, + remember_sync, + remember_in_state, + Rememberer().remember_method, + ) + } + + with subtests.test("injected parameters are absent from the schema"): + for name, tool in tools.items(): + params = tool.info.parameters + assert params is not None + # tool_call_id is hidden just like state, so the LLM only sees x + assert set(params.properties) == {"x"}, name + assert params.required == ["x"], name + + # state is passed for every call below, exercising the filtering of it out of + # the signatures that don't declare it + await dummy_env.reset() + for name, tool in tools.items(): + with subtests.test(f"injected into {name}"): + dummy_env.tools = [tool] + action = ToolRequestMessage( + tool_calls=[ToolCall.from_name(name, id=f"{name}-id", x=1)] + ) + (response,) = await dummy_env.exec_tool_calls( + action, state=dummy_env.state + ) + assert response.content == f"{name}-id:1" + # remember_in_state ran above, so state arrived alongside tool_call_id + assert dummy_env.state.reward == 1 + + dummy_env.tools = [tools["remember"]] + + with subtests.test("concurrent calls each see their own id"): + action = ToolRequestMessage( + tool_calls=[ + ToolCall.from_name("remember", id="first", x=1), + ToolCall.from_name("remember", id="second", x=2), + ] + ) + responses = await dummy_env.exec_tool_calls(action, concurrency=True) + assert [r.content for r in responses] == ["first:1", "second:2"] + + with subtests.test("injected id beats an LLM-emitted one"): + action = ToolRequestMessage( + tool_calls=[ + ToolCall.from_name( + "remember", id="authoritative", x=1, tool_call_id="spoofed" + ) + ] + ) + (response,) = await dummy_env.exec_tool_calls(action) + assert response.content == "authoritative:1" + @pytest.mark.asyncio async def test_tool_timing(self) -> None: sleep_time = 0.1