From a920b7eb44ee6e948b84ea059ebb1bdb1f6b8443 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 29 Oct 2026 14:43:26 -0700 Subject: [PATCH] fix(training): TRL tools carry their MCP input schema, nested types included TRL turns an environment_factory's methods into tool schemas with transformers' get_json_schema, which reads only type hints and docstrings, so an argument like an array of objects reached the model as items: {type: object}, its fields, required lists, enums and bounds only in prose. Each environment instance now sets a SchemaTool over each tool method but end_turn: a dict holding the tool's published JSON Schema, which transformers renders as given, that TRL calls by name. The factory checks at build time that transformers renders the tools as given, and transformers is now declared by the trl extra, since the module imports it. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 2 +- src/agent_env/training/_tools.py | 25 +++++- src/agent_env/training/trl.py | 25 ++++-- tst/data/training/qwen3_chat_template.jinja | 89 +++++++++++++++++++++ tst/trl/grpo_test.py | 17 ++-- tst/unit/openenv/environment_test.py | 11 +-- tst/unit/training/tools_test.py | 57 ++++++++++++- tst/unit/training/trl_test.py | 54 +++++++++++++ uv.lock | 3 +- 9 files changed, 259 insertions(+), 24 deletions(-) create mode 100644 tst/data/training/qwen3_chat_template.jinja diff --git a/pyproject.toml b/pyproject.toml index e9ccc83ca..e2eddfb7c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,7 @@ openenv = [ trl = [ "agentenv-framework[openenv]", "datasets>=3", + "transformers>=5.2", # renders the tools as TRL shows them to a model; 5.2 is what TRL's environment_factory needs ] # The AWS store backends (S3, Secrets Manager, DynamoDB, ECR), selected by impl pointer in config.toml. aws = [ @@ -106,7 +107,6 @@ boat = [ dev = [ "agentenv-framework[aws,boat,boxd,daytona,gcp,openenv,sail,tensorlake,tenki,trl,vercel]", # exercise portable optional providers and stores in CI "fastmcp>=3.4,<4", # the standalone-fastmcp mount path of the protocol; <4 keeps mcp 1.x - "transformers>=5.2", # get_json_schema, which reads agent_env.training's tool methods as TRL shows them to a model "Pillow>=10.4", # the protocol's computer extra, whose tests the unit tier runs "av>=17", # the protocol's computer extra, whose tests the unit tier runs and which the recording tests decode with "griffelib==2.3.0", # the plugin API check; exact, since what it reports as a break changes between releases diff --git a/src/agent_env/training/_tools.py b/src/agent_env/training/_tools.py index 6df1359e4..3eca70dae 100644 --- a/src/agent_env/training/_tools.py +++ b/src/agent_env/training/_tools.py @@ -1,6 +1,6 @@ """Python methods for an env's tools, typed and documented from each tool's JSON Schema the way transformers' -``get_json_schema`` reads a function, which is how TRL shows a model an environment's tools. Built without generating -source; a method calls its instance's ``_call(name, arguments)``.""" +``get_json_schema`` reads a function, which is how TRL finds an environment's tools, and the ``SchemaTool`` each one +reaches the model as. Built without generating source; a method calls its instance's ``_call(name, arguments)``.""" from __future__ import annotations @@ -47,13 +47,32 @@ def method(self: Any, **arguments: Any) -> str: return self._call(name, arguments) method.__name__ = method.__qualname__ = name - method.__doc__ = (tool.get("description") or f"Call the {name} tool.") + ( + method.__doc__ = _summary(tool) + ( "\n\nArgs:\n" + "\n".join(described) if described else "") method.__annotations__ = annotations method.__signature__ = inspect.Signature(parameters, return_annotation=str) return method +class SchemaTool(dict): + """A tool as both its env's JSON Schema, which transformers renders as given, and the callable TRL calls: a method's + type hints have no nested types. TRL lists the methods on the class and takes the tools from the instance, which + sets one over each; a tool has no docstring, so ``get_json_schema`` refuses it rather than make up a schema.""" + + def __init__(self, tool: Mapping[str, Any], call: Callable[..., str]): + super().__init__(type="function", function={ + "name": tool["name"], "description": _summary(tool), "parameters": tool["input_schema"]}) + self.__name__, self.__doc__ = tool["name"], None + self._function = call + + def __call__(self, **arguments: Any) -> str: + return self._function(**arguments) + + +def _summary(tool: Mapping[str, Any]) -> str: + return tool.get("description") or f"Call the {tool['name']} tool." + + def _callable_name(name: str) -> bool: return name.isidentifier() and not keyword.iskeyword(name) and not name.startswith("_") diff --git a/src/agent_env/training/trl.py b/src/agent_env/training/trl.py index 21710f739..d84747ce5 100644 --- a/src/agent_env/training/trl.py +++ b/src/agent_env/training/trl.py @@ -1,7 +1,7 @@ """TRL's trainers on a bundle served as an OpenEnv env (``agent_env.openenv.build_app``, or the package ``openenv import`` writes): ``AgentEnv`` reads the server's tasks into a dataset, and makes the ``environment_factory`` whose environments each play one episode of a task over OpenEnv's WebSocket, offer its turn's tools as methods and own its reward. -``GRPOTrainer`` takes both as they are, verified with TRL 1.16, and AsyncGRPO takes the same ``environment_factory``; +``GRPOTrainer`` takes both as they are, verified with TRL 1.16 and 1.17, and AsyncGRPO the same ``environment_factory``; either needs TRL 1.15 or later, for an environment that owns its reward. Every session runs on one event-loop thread of this module's own, since TRL calls an environment synchronously, from its own threads.""" @@ -29,8 +29,9 @@ try: from datasets import Dataset from openenv.core.generic_client import GenericEnvClient + from transformers.utils.chat_template_utils import render_jinja_template except ModuleNotFoundError as exc: - if exc.name not in ("datasets", "openenv"): + if exc.name not in ("datasets", "openenv", "transformers"): raise raise ImportError("agent_env.training.trl requires pip install 'agentenv-framework[trl]'") from exc @@ -38,7 +39,7 @@ from agent_env.openenv.environment import END_TURN as END_TURN_TOOL from agent_env.openenv.environment import TIMEOUT_METADATA, Toolset from agent_env.task.teardown import TERMINATE_TIMEOUT_SECONDS -from agent_env.training._tools import tool_method +from agent_env.training._tools import SchemaTool, tool_method logger = logging.getLogger(__name__) @@ -247,8 +248,9 @@ class _Rollout: class _Environment: """One session on the server, playing one episode at a time. Its public methods other than ``reset`` and - ``get_reward`` are the tools, which ``AgentEnv.factory`` binds per toolset; every call on the session waits its - turn, as AsyncGRPO can call two tools of a completion at once.""" + ``get_reward`` are the tools, which ``AgentEnv.factory`` binds per toolset, and each but ``end_turn`` is set over + with its ``SchemaTool``; every call on the session waits its turn, as AsyncGRPO can call two tools of a completion + at once.""" _shared: ClassVar[_Shared] _tools: ClassVar[tuple[dict[str, Any], ...]] @@ -258,6 +260,9 @@ def __init__(self) -> None: self._opened = False self._rollout: _Rollout | None = None self._lock = asyncio.Lock() + for tool in self._tools: + if tool["name"] != END_TURN: + setattr(self, tool["name"], SchemaTool(tool, getattr(self, tool["name"]))) _environments.add(self) def reset(self, *, split: str, index: int, task_id: str | None = None, **_: Any) -> None: @@ -413,9 +418,19 @@ def _environment_class(name: str, tools: tuple[dict[str, Any], ...], shared: _Sh methods = {tool["name"]: tool_method(tool) for tool in tools} for method in methods.values(): method.__qualname__ = f"{name}.{method.__name__}" + _check_rendered([SchemaTool(tool, methods[tool["name"]]) for tool in tools if tool["name"] != END_TURN]) return _EnvironmentClass(name, (_Environment,), {**methods, "_shared": shared, "_tools": tools}) +def _check_rendered(tools: list[SchemaTool]) -> None: + """Fails as the trainer sets up, not as the model meets the tools, when transformers doesn't render a tool given + as its JSON Schema as it is.""" + [rendered], _ = render_jinja_template([[]], tools=tools, chat_template="{{ tools | tojson }}") + if changed := [tool.__name__ for tool, shown in zip(tools, json.loads(rendered), strict=True) if shown != tool]: + raise RuntimeError(f"this transformers renders the tools {changed} other than as their env's JSON Schema, so " + "TRL would show the model other tools than the env's") + + def _reduce_environment_class(environment: _EnvironmentClass) -> tuple[Any, ...] | str: """A class ``_environment_class`` built pickles as what it's built from, and a subclass of one by its name, as any class does, so it keeps what it overrides.""" diff --git a/tst/data/training/qwen3_chat_template.jinja b/tst/data/training/qwen3_chat_template.jinja new file mode 100644 index 000000000..01be9b307 --- /dev/null +++ b/tst/data/training/qwen3_chat_template.jinja @@ -0,0 +1,89 @@ +{%- if tools %} + {{- '<|im_start|>system\n' }} + {%- if messages[0].role == 'system' %} + {{- messages[0].content + '\n\n' }} + {%- endif %} + {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} +{%- else %} + {%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('') and message.content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} +{%- endfor %} +{%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if loop.index0 > ns.last_query_index %} + {%- if loop.last or (not loop.last and reasoning_content) %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls %} + {%- for tool_call in message.tool_calls %} + {%- if (loop.first and content) or (not loop.first) %} + {{- '\n' }} + {%- endif %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- endif %} +{%- endif %} \ No newline at end of file diff --git a/tst/trl/grpo_test.py b/tst/trl/grpo_test.py index b5c210f8d..3c9b9a12f 100644 --- a/tst/trl/grpo_test.py +++ b/tst/trl/grpo_test.py @@ -1,6 +1,7 @@ """TRL's GRPOTrainer trains a step through ``agent_env.training.trl`` on CPU, against the OpenEnv server ``build_app`` makes over the Episode API's fake steps. TRL marks ``environment_factory`` experimental, so this pins what the helper -relies on: TRL resets an environment with its row, calls its tools, reads ``get_reward``, and trains on the rewards. +relies on: TRL shows the model each tool as its env publishes it, resets an environment with its row, calls its tools, +reads ``get_reward``, and trains on the rewards. As in TRL's own environment tests, a tiny model's ``generate`` is stood in for. CI runs it against TRL 1.16, the oldest that trains on CPU, and the latest; AsyncGRPO needs a vLLM server, so it isn't covered here.""" @@ -8,7 +9,8 @@ import pytest -from tst.unit.training.trl_test import PLAYED +from tst.unit.training.tools_test import tools_in +from tst.unit.training.trl_test import PLAYED, SHOWN, STACK from tst.util.capabilities import TRL, missing_capability_reason torch = pytest.importorskip("torch", reason=missing_capability_reason(TRL)) @@ -19,9 +21,10 @@ END = '\n{"name": "end_turn", "arguments": {"response": "Added a crate."}}\n' -def test_grpotrainer_trains_a_step_on_a_served_bundle_and_scores_every_rollout(served, teardowns, tmp_path, +def test_grpotrainer_trains_a_step_on_a_served_bundle_and_scores_every_rollout(served, teardowns, tools, tmp_path, monkeypatch): monkeypatch.setenv("TRL_EXPERIMENTAL_SILENCE", "1") + tools("add_item", STACK) env = served({"a": PLAYED, "b": PLAYED}) rewards = [] @@ -36,10 +39,12 @@ def get_reward(self): use_cpu=True)) before = {name: param.clone() for name, param in trainer.model.named_parameters()} - with patch.object(trainer.model, "generate", side_effect=_Policy(trainer.processing_class)): + policy = _Policy(trainer.processing_class) + with patch.object(trainer.model, "generate", side_effect=policy): trainer.train() [logged] = [row for row in trainer.state.log_history if "tools/call_frequency" in row] + assert tools_in(policy.prompt) == SHOWN assert sorted(rewards) == [0.0, 0.0, 1.0, 1.0] assert (logged["tools/call_frequency"], logged["tools/failure_frequency"], logged["rewards/Recorded/mean"]) == ( pytest.approx(1.0), 0.0, pytest.approx(0.5)) @@ -50,14 +55,16 @@ def get_reward(self): class _Policy: """Stands in for ``generate``: of each group, the first rollout adds a crate, then ends its turn, then stops, and the - second answers without a tool, so its turn is ended empty and scores 0.""" + second answers without a tool, so its turn is ended empty and scores 0. ``prompt`` is the first it was given.""" def __init__(self, tokenizer): self.tokenizer = tokenizer self.calls = 0 + self.prompt = None def __call__(self, input_ids, **_): self.calls += 1 + self.prompt = self.prompt or self.tokenizer.decode(input_ids[0]) if self.calls == 1: texts = [ADD if row % 2 == 0 else "I won't." for row in range(len(input_ids))] else: diff --git a/tst/unit/openenv/environment_test.py b/tst/unit/openenv/environment_test.py index b8d9c5d29..046800be3 100644 --- a/tst/unit/openenv/environment_test.py +++ b/tst/unit/openenv/environment_test.py @@ -21,9 +21,9 @@ class _Tools: - """Stands in for the proxy's ``ToolSession``: every env offers ``offered``, and a call answers with what ``answers`` maps - its tool to, raising it when it's an exception. A session opens once ``release``, a ``threading.Event`` since the - in-process proxy runs on a loop of its own, is set, when there is one.""" + """Stands in for the proxy's ``ToolSession``: every env offers ``offered``, tools or their names, and a call answers + with what ``answers`` maps its tool to, raising it when it's an exception. A session opens once ``release``, a + ``threading.Event`` since the in-process proxy runs on a loop of its own, is set, when there is one.""" def __init__(self, offered, answers): self.offered, self.answers, self.closed, self.release = offered, answers, 0, None @@ -38,8 +38,9 @@ async def open(self): return self async def list_tools(self): - return [EnvironmentTool(name=name, description=f"{name}.", inputSchema={"type": "object"}) - for name in tools.offered] + return [tool if isinstance(tool, EnvironmentTool) + else EnvironmentTool(name=tool, description=f"{tool}.", inputSchema={"type": "object"}) + for tool in tools.offered] async def call_tool(self, name, arguments=None): answer = tools.answers.get(name, ToolResult(content=[{"type": "text", "text": f"{name} {arguments}"}])) diff --git a/tst/unit/training/tools_test.py b/tst/unit/training/tools_test.py index f65e00d59..a7c395f01 100644 --- a/tst/unit/training/tools_test.py +++ b/tst/unit/training/tools_test.py @@ -1,13 +1,30 @@ -"""The tool methods TRL calls an env's tools through: their signatures and docstrings from each tool's JSON Schema, the -schema transformers' ``get_json_schema`` reads back from them, which is what TRL shows a model, and the tools refused.""" +"""The tool methods TRL finds an env's tools by: their signatures and docstrings from each tool's JSON Schema, the +schema transformers' ``get_json_schema`` reads back from them, and the tools refused; and the ``SchemaTool`` that shows +a model the tool as its env publishes it, rendered by Qwen3's chat template and parsed back from its tool call.""" import inspect +import json +from pathlib import Path from typing import Any, Literal import pytest -from transformers.utils import get_json_schema +from tokenizers import Tokenizer, models +from transformers import PreTrainedTokenizerFast +from transformers.utils import DocstringParsingException, get_json_schema +from transformers.utils.chat_parsing import parse_response -from agent_env.training._tools import tool_method +from agent_env.training._tools import SchemaTool, tool_method + +QWEN3 = PreTrainedTokenizerFast( + tokenizer_object=Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")), + chat_template=(Path(__file__).parents[2] / "data" / "training" / "qwen3_chat_template.jinja").read_text()) +#: Qwen3's response template as TRL sets it on the tokenizer, which ``parse_response`` reads a completion with. +QWEN3_RESPONSE = {"defaults": {"role": "assistant"}, "start_anchor": "<|im_start|>assistant\n", "fields": { + "reasoning_content": {"open": "", "close_pattern": r"\s*", "content": "text"}, + "tool_calls": {"open": "", "close_pattern": r"\s*", "repeats": True, "content": "json", + "transform": {"type": "function", "function": "{content}"}}, + "content": {"close_pattern": r"<\|im_end\|>\s*", "content": "text"}, +}} RESTOCK = {"name": "restock", "description": "Restock a shelf.", "input_schema": { "type": "object", "required": ["shelf", "count"], @@ -82,3 +99,35 @@ def test_transformers_reads_back_each_tool_as_the_env_describes_it(): def test_a_tool_trl_cant_show_the_model_is_refused(tool: dict[str, Any], message): with pytest.raises(ValueError, match=message): tool_method(tool) + + +def test_a_schema_tool_shows_the_model_its_envs_schema_nested_types_included_and_calls_through(): + called = [] + tool = SchemaTool(PLAN, lambda **arguments: called.append(arguments) or "planned") + + assert shown_to_qwen3([tool]) == [{"type": "function", "function": { + "name": "check_plan", "description": "Check a plan.", "parameters": PLAN["input_schema"]}}] + [call] = called_by_qwen3([tool], {"name": "check_plan", "arguments": {"plan": json.dumps([{"ship": 3}])}}) + assert call == {"type": "function", "function": {"name": "check_plan", "arguments": {"plan": [{"ship": 3}]}}} + assert tool(**call["function"]["arguments"]) == "planned" + assert called == [{"plan": [{"ship": 3}]}] + with pytest.raises(DocstringParsingException, match="for check_plan because it has no docstring"): + get_json_schema(tool) + + +def shown_to_qwen3(tools: list[Any]) -> list[dict[str, Any]]: + """The tools Qwen3's chat template shows the model, as TRL renders a prompt.""" + return tools_in(QWEN3.apply_chat_template([{"role": "user", "content": "Plan."}], tools=tools, tokenize=False)) + + +def tools_in(prompt: str) -> list[dict[str, Any]]: + return [json.loads(line) for line in prompt.split("\n")[1].split("\n")[0].splitlines()] + + +def called_by_qwen3(tools: list[Any], *calls: dict[str, Any]) -> list[dict[str, Any]]: + """The tool calls ``parse_response`` reads from Qwen3's completion making ``calls``, typed by ``tools``' schemas as + TRL parses a completion.""" + prompt = QWEN3.apply_chat_template([{"role": "user", "content": "Plan."}], tools=tools, tokenize=False, + add_generation_prompt=True) + completion = "\n\n\n\n" + "".join(f"\n{json.dumps(call)}\n" for call in calls) + return parse_response(completion + "<|im_end|>", QWEN3_RESPONSE, prefix=prompt, tools=tools)["tool_calls"] diff --git a/tst/unit/training/trl_test.py b/tst/unit/training/trl_test.py index c9a2cc9dc..17fc6d85b 100644 --- a/tst/unit/training/trl_test.py +++ b/tst/unit/training/trl_test.py @@ -24,6 +24,7 @@ from tst.unit.bundle._support import layout from tst.unit.openenv.environment_test import AGAIN from tst.unit.task.episode.run_test import ALICE, ITEMS, PANTRY, PANTRY_MCP, REPLIES, TURN, _tool_sessions +from tst.unit.training.tools_test import called_by_qwen3, shown_to_qwen3 PLAYED = [ITEMS, ALICE, TURN, REPLIES] SYSTEM = {"role": "system", "content": "You keep the items."} @@ -327,6 +328,59 @@ def test_tasks_whose_agent_gets_other_tools_take_trls_multi_environment_mode(ser assert pantry.get_reward() == 1.0 +STACK = EnvironmentTool(name="stack_crates", description="Stack crates on a pallet.", inputSchema={ + "type": "object", "required": ["pallet", "crates"], "properties": { + "pallet": {"type": "string", "enum": ["north", "south"]}, + "crates": {"type": "array", "description": "bottom crate first", "items": { + "type": "object", "required": ["kind"], "properties": { + "kind": {"type": "string", "description": "what the crate holds"}, + "size": {"type": "array", "items": {"type": "number"}, "minItems": 3, "maxItems": 3}}}}}}) +#: What TRL shows the model of the tools ``add_item`` and ``STACK``: each as its env publishes it, and ``end_turn`` as +#: TRL has shown it since the factory's first release, its method's schema. +SHOWN = [ + {"type": "function", "function": {"name": "add_item", "description": "add_item.", "parameters": {"type": "object"}}}, + {"type": "function", "function": { + "name": "end_turn", + "description": "End your turn with your reply. The episode goes on to your next turn, or ends and is scored.", + "parameters": {"type": "object", "properties": { + "response": {"type": "string", "nullable": True, "description": "Your reply."}}}, + "return": {"type": "string"}}}, + {"type": "function", "function": {"name": "stack_crates", "description": "Stack crates on a pallet.", + "parameters": STACK.inputSchema}}] + + +def test_trl_shows_the_model_each_tool_as_its_env_publishes_it_and_end_turn_as_before_and_calls_them(served, tools): + tools("add_item", STACK) + env = served({"a": PLAYED}) + environment = env.factory() + fetched = [getattr(environment, name) for name in _tools(env.factory) if name not in ("reset", "get_reward")] + + assert shown_to_qwen3(fetched) == SHOWN + from_the_class = env.factory.stack_crates.__get__(environment) + assert shown_to_qwen3([from_the_class])[0]["function"]["parameters"]["properties"]["crates"]["items"] == { + "type": "object"} + environment.reset(**env.dataset[0]) + crates = [{"kind": "apples", "size": [1, 1, 0.5]}] + calls = called_by_qwen3(fetched, {"name": "stack_crates", "arguments": {"pallet": "north", "crates": crates}}, + {"name": "end_turn", "arguments": {"response": "Added a crate."}}) + by_name = {tool.__name__: tool for tool in fetched} + + assert [by_name[call["function"]["name"]](**call["function"]["arguments"]) for call in calls] == [ + f"stack_crates {{'pallet': 'north', 'crates': {crates}}}", ENDED] + assert environment.get_reward() == 1.0 + + +def test_a_transformers_that_renders_a_tool_other_than_as_its_schema_stops_the_factory(served, tools, monkeypatch): + tools(STACK) + env = served({"a": PLAYED}) + monkeypatch.setattr(trl, "render_jinja_template", lambda conversations, tools, chat_template: ( + [json.dumps([{**tool, "function": {**tool["function"], "parameters": {"type": "object"}}} for tool in tools])], + [])) + + with pytest.raises(RuntimeError, match=r"renders the tools \['stack_crates'\] other than as their env's"): + env.factory # noqa: B018 + + def test_a_tool_trl_cant_call_is_refused_with_how_to_fix_it(served, tools): tools("add-item") env = served({"a": PLAYED}) diff --git a/uv.lock b/uv.lock index 0b9b4091d..1faf5c1b7 100644 --- a/uv.lock +++ b/uv.lock @@ -145,6 +145,7 @@ tensorlake = [ trl = [ { name = "datasets" }, { name = "openenv" }, + { name = "transformers" }, ] vercel = [ { name = "anyio" }, @@ -206,7 +207,7 @@ requires-dist = [ { name = "tenki", extras = ["async"], marker = "extra == 'tenki'", specifier = "~=1.4.1" }, { name = "tensorlake", marker = "extra == 'tensorlake'", specifier = "~=0.5.140" }, { name = "tomli-w", marker = "extra == 'dev'", specifier = "==1.2.0" }, - { name = "transformers", marker = "extra == 'dev'", specifier = ">=5.2" }, + { name = "transformers", marker = "extra == 'trl'", specifier = ">=5.2" }, { name = "urllib3", marker = "extra == 'boat'", specifier = ">=2.1.0,<3" }, { name = "uvicorn", specifier = ">=0.47" }, { name = "vercel-sandbox", marker = "extra == 'vercel'", specifier = ">=0.7.0,<0.8" },