Send back helpful message instead of raising when action is plain Message

packages/gsm8k/src/aviary/envs/gsm8k/env.py

@@ -15,6 +15,7 @@
     Messages,
     TaskDataset,
     Tool,
+    ToolRequestMessage,
     ToolResponseMessage,
 )
 from pydantic import BaseModel, ConfigDict
@@ -141,7 +142,8 @@ async def reset(self) -> tuple[Messages, list[Tool]]:
         return [Message(content=self.problem)], self.tools
 
     async def step(self, action: Message) -> tuple[Messages, float, bool, bool]:
-        action = self.check_action_is_tool_request(action)
+        if not isinstance(action, ToolRequestMessage):
+            return self.default_no_tool_calls_response
         if not action.tool_calls:
             return (
                 [

packages/hotpotqa/src/aviary/envs/hotpotqa/env.py

@@ -33,6 +33,7 @@
     Messages,
     TaskDataset,
     Tool,
+    ToolRequestMessage,
     eval_answer,
 )
 from bs4 import BeautifulSoup
@@ -300,7 +301,8 @@ async def step(self, action: Message) -> tuple[Messages, float, bool, bool]:
             >>> print(obs, done)
             [ToolResponseMessage(name='Search', tool_call_id='tool_call_id', content='...')], False
         """
-        action = self.check_action_is_tool_request(action)
+        if not isinstance(action, ToolRequestMessage):
+            return self.default_no_tool_calls_response
         self.state.steps += 1
         if not action.tool_calls:
             return (

packages/labbench/src/aviary/envs/labbench/env.py

@@ -12,6 +12,7 @@
     Messages,
     MultipleChoiceEvaluation,
     MultipleChoiceQuestion,
+    ToolRequestMessage,
 )
 from aviary.env import ENV_REGISTRY
 from ldp.utils import discounted_returns
@@ -141,8 +142,9 @@ async def _evaluate_answer(self) -> TEvaluation:
         return evaluation  # type: ignore[return-value]
 
     async def step(self, action: Message) -> tuple[Messages, float, bool, bool]:
-        tool_request = self.check_action_is_tool_request(action)
-        messages, reward, done, truncated = await super().step(tool_request)
+        if not isinstance(action, ToolRequestMessage):
+            return self.default_no_tool_calls_response
+        messages, reward, done, truncated = await super().step(action)
         if not done or not isinstance(self._query, MultipleChoiceQuestion):
             return messages, reward, done, truncated
         evaluation = await self._evaluate_answer()

packages/lfrqa/src/aviary/envs/lfrqa/env.py

@@ -14,6 +14,7 @@
     Message,
     Messages,
     MultipleChoiceQuestion,
+    ToolRequestMessage,
 )
 from aviary.env import ENV_REGISTRY
 from aviary.envs.labbench import GradablePaperQAEnvironment
@@ -270,10 +271,11 @@ async def _evaluate_answer(self) -> dict:
         return evaluation
 
     async def step(self, action: Message) -> tuple[Messages, float, bool, bool]:
-        tool_request = self.check_action_is_tool_request(action)
+        if not isinstance(action, ToolRequestMessage):
+            return self.default_no_tool_calls_response
         messages, reward, done, truncated = await super(
             GradablePaperQAEnvironment, self
-        ).step(tool_request)
+        ).step(action)
         if not done:
             return messages, reward, done, truncated
         evaluation = await self._evaluate_answer()

packages/notebook/src/aviary/envs/notebook/env.py

@@ -11,7 +11,7 @@
 
 import aiodocker
 import nbformat
-from aviary.core import Environment, Message, Messages, Tool
+from aviary.core import Environment, Message, Messages, Tool, ToolRequestMessage
 from aviary.message import EnvStateMessage
 from jupyter_client.manager import AsyncKernelManager
 from nbformat import NotebookNode
@@ -184,14 +184,13 @@ async def reset(self) -> tuple[Messages, list[Tool]]:
         return init_obs, self.tools
 
     async def step(self, action: Message) -> tuple[Messages, float, bool, bool]:
-        tool_request = self.check_action_is_tool_request(action)
+        if not isinstance(action, ToolRequestMessage):
+            return self.default_no_tool_calls_response
         prev_reward = self.state.total_reward
 
         obs = cast(
             Messages,
-            await self.exec_tool_calls(
-                tool_request, concurrency=False, handle_tool_exc=True
-            ),
+            await self.exec_tool_calls(action, concurrency=False, handle_tool_exc=True),
         )
         reward = self.state.total_reward - prev_reward
 

src/aviary/env.py

@@ -103,13 +103,18 @@ class Environment(ABC, Generic[TEnvState]):
     tools: list[Tool]
     state: TEnvState
 
-    @staticmethod
-    def check_action_is_tool_request(action: Message) -> ToolRequestMessage:
-        if not isinstance(action, ToolRequestMessage):
-            raise TypeError(
-                f"Expected a ToolRequestMessage, but got {type(action).__name__}."
-            )
-        return action
+    @property
+    def default_no_tool_calls_response(self) -> tuple[Messages, float, bool, bool]:
+        return (
+            [
+                Message(
+                    content="No tool calls received. Please call one or more tools to proceed."
+                )
+            ],
+            0.0,
+            False,
+            False,
+        )
 
     @abstractmethod
     async def step(self, action: Message) -> tuple[Messages, float, bool, bool]:
@@ -552,9 +557,10 @@ def from_task(cls, task: str) -> "DummyEnv":
         return cls(task=task)
 
     async def step(self, action: Message) -> tuple[Messages, float, bool, bool]:
-        tool_request = self.check_action_is_tool_request(action)
+        if not isinstance(action, ToolRequestMessage):
+            return self.default_no_tool_calls_response
         msgs: Messages = await self.exec_tool_calls(  # type: ignore[assignment]
-            tool_request, state=self.state, concurrency=self.concurrent_tool_calls
+            action, state=self.state, concurrency=self.concurrent_tool_calls
         ) or [
             ToolResponseMessage(
                 content=f"No tool calls input in tool request {action}.",

src/aviary/functional.py

@@ -6,7 +6,7 @@
 
 from aviary.env import Environment, Frame
 from aviary.message import Message
-from aviary.tools import Messages, Tool
+from aviary.tools import Messages, Tool, ToolRequestMessage
 from aviary.utils import is_coroutine_callable
 
 
@@ -72,9 +72,10 @@ async def reset(self) -> tuple[Messages, list[Tool]]:
         return [Message(content=obs)], self.tools
 
     async def step(self, action: Message) -> tuple[Messages, float, bool, bool]:
-        tool_request = self.check_action_is_tool_request(action)
+        if not isinstance(action, ToolRequestMessage):
+            return self.default_no_tool_calls_response
         msgs = await self.exec_tool_calls(
-            tool_request,
+            action,
             state=self.state,
             concurrency=self.allow_concurrency,
             handle_tool_exc=True,

tests/test_envs.py

@@ -353,8 +353,9 @@ def slow_tool() -> None:
         return [], self.tools
 
     async def step(self, action: Message) -> tuple[Messages, float, bool, bool]:
-        tool_request = self.check_action_is_tool_request(action)
-        await self.exec_tool_calls(tool_request, exec_timeout=0.0001)
+        if not isinstance(action, ToolRequestMessage):
+            return self.default_no_tool_calls_response
+        await self.exec_tool_calls(action, exec_timeout=0.0001)
 
         return [], 0.0, False, False
 
@@ -754,9 +755,12 @@ async def test_step_with_plain_message(self, server_async_client: AsyncClient):
         response = await server_async_client.post(
             "/step", json={"env_id": env_id, "action": action.model_dump()}
         )
-        assert response.status_code == 500, (
-            "Plain Message should fail check_action_is_tool_request"
-        )
+        assert response.status_code == 200
+        obs, reward, done, truncated = response.json()
+        assert any("No tool calls" in m["content"] for m in obs)
+        assert reward == 0.0
+        assert not done
+        assert not truncated
 
     @pytest.mark.asyncio
     async def test_step_with_tool_response_message(
@@ -772,9 +776,34 @@ async def test_step_with_tool_response_message(
         response = await server_async_client.post(
             "/step", json={"env_id": env_id, "action": action.model_dump()}
         )
-        assert response.status_code == 500, (
-            "ToolResponseMessage should fail check_action_is_tool_request"
-        )
+        assert response.status_code == 200
+        obs, reward, done, truncated = response.json()
+        assert any("No tool calls" in m["content"] for m in obs)
+        assert reward == 0.0
+        assert not done
+        assert not truncated
+
+
+class TestDefaultNoToolCallsResponse:
+    @pytest.mark.parametrize(
+        "action",
+        [
+            Message(content="hello"),
+            ToolResponseMessage(content="result", name="tool", tool_call_id="abc"),
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_non_tool_request_returns_default(
+        self, dummy_env: DummyEnv, action: Message
+    ):
+        await dummy_env.reset()
+        obs, reward, done, truncated = await dummy_env.step(action)
+        assert len(obs) == 1
+        assert obs[0].content
+        assert "No tool calls" in obs[0].content
+        assert reward == 0.0
+        assert not done
+        assert not truncated
 
 
 class TestStepRequestDeserialization: