Repository: Future-House/aviary

Issue:
Move `eval_answer` from `tools`
Should not live in `tools`, since it's more of an environment function

Discussion:
maykcaldas: However, `aviary_internal` uses this function in two envs

Patch:
diff --git a/src/aviary/core.py b/src/aviary/core.py
index 8c3369d..587e76b 100644
--- a/src/aviary/core.py
+++ b/src/aviary/core.py
@@ -20,7 +20,6 @@ from aviary.message import EnvStateMessage, MalformedMessageError, Message, join
 from aviary.render import Renderer
 from aviary.tools import (
     INVALID_TOOL_NAME,
-    EvalAnswerMode,
     FunctionInfo,
     Messages,
     MessagesAdapter,
@@ -35,10 +34,15 @@ from aviary.tools import (
     ToolSelector,
     ToolSelectorLedger,
     argref_by_name,
-    eval_answer,
     wraps_doc_only,
 )
-from aviary.utils import encode_image_to_base64, is_coroutine_callable, partial_format
+from aviary.utils import (
+    EvalAnswerMode,
+    encode_image_to_base64,
+    eval_answer,
+    is_coroutine_callable,
+    partial_format,
+)
 
 __all__ = [
     "INVALID_TOOL_NAME",
@@ -77,6 +81,7 @@ __all__ = [
     "argref_by_name",
     "encode_image_to_base64",
     "eval_answer",
+    "eval_answer",
     "fenv",
     "is_coroutine_callable",
     "join",
diff --git a/src/aviary/env.py b/src/aviary/env.py
index e38edb3..b8be875 100644
--- a/src/aviary/env.py
+++ b/src/aviary/env.py
@@ -7,9 +7,9 @@ import json
 import logging
 import random
 from abc import ABC, abstractmethod
-from collections.abc import Iterator
+from collections.abc import Awaitable, Iterator
 from copy import deepcopy
-from typing import Annotated, Generic, Self, TypeAlias, TypeVar, cast
+from typing import Annotated, Any, Generic, Self, TypeAlias, TypeVar, cast
 
 from pydantic import (
     BaseModel,
@@ -39,6 +39,13 @@ logger = logging.getLogger(__name__)
 Serializable: TypeAlias = dict | list | int | float | str | bool | BaseModel
 
 
+async def maybe_wait_for(future: Awaitable, timeout: float | None) -> Any:
+    """Apply a timeout to an awaitable if one is provided, otherwise await indefinitely."""
+    if timeout is None:
+        return await future
+    return await asyncio.wait_for(future, timeout)
+
+
 class Frame(BaseModel):
     """A frame is a snapshot at a given timestep. The name comes from video frame."""
 
@@ -158,6 +165,7 @@ class Environment(ABC, Generic[TEnvState]):
         concurrency: bool = True,
         handle_tool_exc: bool = False,
         handle_invalid_tool_calls: bool = True,
+        exec_timeout: float | None = None,
         **function_kwargs,
     ) -> list[ToolResponseMessage]:
         """
@@ -171,6 +179,8 @@ class Environment(ABC, Generic[TEnvState]):
                 ToolResponseMessage.
             handle_invalid_tool_calls: Flag to handle invalid tool calls by returning
                 a ToolResponseMessage with a note that the tool requested doesn't exist
+            exec_timeout: Timeout for each tool call in seconds. If None, no timeout.
+                Note that handle_tool_exec can be used to catch TimeoutErrors.
             **function_kwargs: Keyword arguments to pass to all tool functions.
 
         Returns:
@@ -203,15 +213,21 @@ class Environment(ABC, Generic[TEnvState]):
             tool_exc: Exception | None = None
             try:
                 if is_coroutine_callable(tool._tool_fn):
-                    content = await tool._tool_fn(
-                        **tool_call.function.arguments, **filtered_kwargs
+                    content = await maybe_wait_for(
+                        tool._tool_fn(
+                            **tool_call.function.arguments, **filtered_kwargs
+                        ),
+                        exec_timeout,
                     )
                 else:
                     # If the function is synchronous, run on a thread
-                    content = await asyncio.to_thread(
-                        tool._tool_fn,
-                        **tool_call.function.arguments,
-                        **filtered_kwargs,
+                    content = await maybe_wait_for(
+                        asyncio.to_thread(
+                            tool._tool_fn,
+                            **tool_call.function.arguments,
+                            **filtered_kwargs,
+                        ),
+                        exec_timeout,
                     )
             except Exception as exc:
                 if not handle_tool_exc:
diff --git a/src/aviary/tools/__init__.py b/src/aviary/tools/__init__.py
index c03b8cd..e18c552 100644
--- a/src/aviary/tools/__init__.py
+++ b/src/aviary/tools/__init__.py
@@ -14,11 +14,10 @@ from .base import (
     ToolsAdapter,
     wraps_doc_only,
 )
-from .utils import EvalAnswerMode, ToolSelector, ToolSelectorLedger, eval_answer
+from .utils import ToolSelector, ToolSelectorLedger
 
 __all__ = [
     "INVALID_TOOL_NAME",
-    "EvalAnswerMode",
     "FunctionInfo",
     "Messages",
     "MessagesAdapter",
@@ -33,6 +32,5 @@ __all__ = [
     "Tools",
     "ToolsAdapter",
     "argref_by_name",
-    "eval_answer",
     "wraps_doc_only",
 ]
diff --git a/src/aviary/tools/utils.py b/src/aviary/tools/utils.py
index e0fe363..36a4987 100644
--- a/src/aviary/tools/utils.py
+++ b/src/aviary/tools/utils.py
@@ -1,5 +1,4 @@
 from collections.abc import Callable
-from enum import StrEnum
 from functools import partial
 from typing import TYPE_CHECKING, Any, ClassVar, cast
 
@@ -21,103 +20,6 @@ if TYPE_CHECKING:
     from litellm import ModelResponse
 
 
-class EvalAnswerMode(StrEnum):
-    EXACT = "exact"  # strings must match exactly
-    CONTAINS = "contains"  # the correct answer is contained in the supplied answer
-    LLM = "llm"  # Ask an LLM to evaluate
-    LLM_SCORE = "llm-score"  # Ask an LLM to evaluate and return the score (normalized)
-
-
-LLM_EVAL_CONFIG = {
-    "prompt": (
-        "Here is a question, the correct answer to the question, and a proposed answer"
-        " to the question. Please tell me if the proposed answer is correct, given the"
-        " correct answer. ONLY SAY 'YES' OR 'NO'. No other output is permitted."
-        "\n\nQuestion: {question}"
-        "\n\nCorrect answer: {correct_answer}"
-        "\n\nProposed answer: {proposed_answer}"
-    ),
-    "model": "gpt-4o-mini",
-    "temperature": 0,
-}
-
-LLM_SCORE_EVAL_CONFIG = {
-    "prompt": (
-        "Here is a question, the correct answer to the question, and a rubric for"
-        " evaluating the question. Judge the proposed answer based on the given rubric."
-        " Give a score from 0 to 10. No other output is permitted."
-        "\n\nQuestion: {question}"
-        "\n\nRubric: {correct_answer}"
-        "\n\nProposed answer: {proposed_answer}"
-    ),
-    "model": "gpt-4o-mini",
-    "temperature": 0,
-    "max_score": 10,
-}
-
-
-async def eval_answer(
-    proposed: str,
-    correct: str,
-    question: str | None = None,
-    eval_mode: EvalAnswerMode = EvalAnswerMode.CONTAINS,
-    llm_eval_config: dict | None = None,
-) -> float:
-    """Evaluate a proposed answer against a correct answer.
-
-    Will return 0 or 1, except for llm-score which should be between 0 and 1
-    """
-    if eval_mode in {EvalAnswerMode.LLM, EvalAnswerMode.LLM_SCORE}:
-        try:
-            from litellm import acompletion
-        except ImportError as e:
-            raise ImportError(
-                "eval_answer requires the 'llm' extra for 'litellm'. Please:"
-                " `pip install aviary[llm]`."
-            ) from e
-        if question is None:
-            raise ValueError("Question must be provided for LLM evaluation mode.")
-        default_config = (
-            LLM_EVAL_CONFIG
-            if eval_mode == EvalAnswerMode.LLM
-            else LLM_SCORE_EVAL_CONFIG
-        )
-        config = llm_eval_config or default_config
-        prompt = cast(str, config.get("prompt", default_config["prompt"])).format(
-            question=question,
-            correct_answer=correct,
-            proposed_answer=proposed,
-        )
-        response = await acompletion(
-            model=config.get("model", default_config["model"]),
-            temperature=config.get("temperature", default_config["temperature"]),
-            messages=[{"content": prompt, "role": "user"}],
-        )
-        if eval_mode == EvalAnswerMode.LLM:
-            return await eval_answer(
-                response.choices[0].message.content.strip().casefold(),
-                "yes",
-                eval_mode=EvalAnswerMode.EXACT,
-            )
-        try:
-            return float(response.choices[0].content.strip()) / float(
-                config.get("max_score", default_config["max_score"])  # type: ignore[arg-type]
-            )
-        except ValueError:
-            return 0
-
-    gt = correct.strip().casefold()
-    pred = proposed.strip().casefold()
-
-    if eval_mode == EvalAnswerMode.EXACT:
-        return float(pred == gt)
-
-    if eval_mode == EvalAnswerMode.CONTAINS:
-        return float(gt in pred)
-
-    raise RuntimeError(f"Invalid evaluation mode: {eval_mode}")
-
-
 class ToolSelector:
     """Simple entity to select a tool based on messages."""
 
diff --git a/src/aviary/utils.py b/src/aviary/utils.py
index e7d280b..4d4fb3e 100644
--- a/src/aviary/utils.py
+++ b/src/aviary/utils.py
@@ -2,12 +2,48 @@ import base64
 import contextlib
 import inspect
 import io
-from typing import TYPE_CHECKING, Any
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any, cast
 
 if TYPE_CHECKING:
     import numpy as np
 
 
+LLM_EVAL_CONFIG = {
+    "prompt": (
+        "Here is a question, the correct answer to the question, and a proposed answer"
+        " to the question. Please tell me if the proposed answer is correct, given the"
+        " correct answer. ONLY SAY 'YES' OR 'NO'. No other output is permitted."
+        "\n\nQuestion: {question}"
+        "\n\nCorrect answer: {correct_answer}"
+        "\n\nProposed answer: {proposed_answer}"
+    ),
+    "model": "gpt-4o-mini",
+    "temperature": 0,
+}
+
+LLM_SCORE_EVAL_CONFIG = {
+    "prompt": (
+        "Here is a question, the correct answer to the question, and a rubric for"
+        " evaluating the question. Judge the proposed answer based on the given rubric."
+        " Give a score from 0 to 10. No other output is permitted."
+        "\n\nQuestion: {question}"
+        "\n\nRubric: {correct_answer}"
+        "\n\nProposed answer: {proposed_answer}"
+    ),
+    "model": "gpt-4o-mini",
+    "temperature": 0,
+    "max_score": 10,
+}
+
+
+class EvalAnswerMode(StrEnum):
+    EXACT = "exact"  # strings must match exactly
+    CONTAINS = "contains"  # the correct answer is contained in the supplied answer
+    LLM = "llm"  # Ask an LLM to evaluate
+    LLM_SCORE = "llm-score"  # Ask an LLM to evaluate and return the score (normalized)
+
+
 def partial_format(value: str, **formats: dict[str, Any]) -> str:
     """Partially format a string given a variable amount of formats."""
     for template_key, template_value in formats.items():
@@ -41,3 +77,65 @@ def is_coroutine_callable(obj) -> bool:
     if callable(obj):
         return inspect.iscoroutinefunction(obj.__call__)
     return False
+
+
+async def eval_answer(
+    proposed: str,
+    correct: str,
+    question: str | None = None,
+    eval_mode: EvalAnswerMode = EvalAnswerMode.CONTAINS,
+    llm_eval_config: dict | None = None,
+) -> float:
+    """Evaluate a proposed answer against a correct answer.
+
+    Will return 0 or 1, except for llm-score which should be between 0 and 1
+    """
+    if eval_mode in {EvalAnswerMode.LLM, EvalAnswerMode.LLM_SCORE}:
+        try:
+            from litellm import acompletion
+        except ImportError as e:
+            raise ImportError(
+                "eval_answer requires the 'llm' extra for 'litellm'. Please:"
+                " `pip install aviary[llm]`."
+            ) from e
+        if question is None:
+            raise ValueError("Question must be provided for LLM evaluation mode.")
+        default_config = (
+            LLM_EVAL_CONFIG
+            if eval_mode == EvalAnswerMode.LLM
+            else LLM_SCORE_EVAL_CONFIG
+        )
+        config = llm_eval_config or default_config
+        prompt = cast(str, config.get("prompt", default_config["prompt"])).format(
+            question=question,
+            correct_answer=correct,
+            proposed_answer=proposed,
+        )
+        response = await acompletion(
+            model=config.get("model", default_config["model"]),
+            temperature=config.get("temperature", default_config["temperature"]),
+            messages=[{"content": prompt, "role": "user"}],
+        )
+        if eval_mode == EvalAnswerMode.LLM:
+            return await eval_answer(
+                response.choices[0].message.content.strip().casefold(),
+                "yes",
+                eval_mode=EvalAnswerMode.EXACT,
+            )
+        try:
+            return float(response.choices[0].content.strip()) / float(
+                config.get("max_score", default_config["max_score"])  # type: ignore[arg-type]
+            )
+        except ValueError:
+            return 0
+
+    gt = correct.strip().casefold()
+    pred = proposed.strip().casefold()
+
+    if eval_mode == EvalAnswerMode.EXACT:
+        return float(pred == gt)
+
+    if eval_mode == EvalAnswerMode.CONTAINS:
+        return float(gt in pred)
+
+    raise RuntimeError(f"Invalid evaluation mode: {eval_mode}")


Test patch:
diff --git a/packages/hotpotqa/tests/test_hotpotqa_env.py b/packages/hotpotqa/tests/test_hotpotqa_env.py
index 9db65c1..f6cdc8d 100644
--- a/packages/hotpotqa/tests/test_hotpotqa_env.py
+++ b/packages/hotpotqa/tests/test_hotpotqa_env.py
@@ -5,7 +5,7 @@ import pytest
 from aviary.core import Environment, TaskDataset
 from aviary.envs.hotpotqa import HotPotQAEnv
 from aviary.envs.hotpotqa.env import HotPotQADataset
-from aviary.tools.utils import EvalAnswerMode
+from aviary.utils import EvalAnswerMode
 
 
 def test_env_construction() -> None:
diff --git a/tests/test_envs.py b/tests/test_envs.py
index 4e3a4e8..4d34945 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -30,6 +30,7 @@ from aviary.core import (
     ToolSelectorLedger,
 )
 from aviary.dataset_server import TaskDatasetServer
+from aviary.tools import Messages
 from tests import CILLMModelNames
 from tests.conftest import VCR_DEFAULT_MATCH_ON
 
@@ -236,6 +237,38 @@ async def test_invalid_tool_call(
         assert o.tool_call_id == t.id
 
 
+class SlowEnv(Environment[None]):
+    async def reset(self) -> tuple[list[Message], list[Tool]]:
+        async def aslow_tool() -> None:
+            """I am very slow."""
+            await asyncio.sleep(0.1)
+
+        def slow_tool() -> None:
+            """I am very slow."""
+            time.sleep(0.1)
+
+        self.tools = [Tool.from_function(slow_tool), Tool.from_function(aslow_tool)]
+        return [], self.tools
+
+    async def step(
+        self, action: ToolRequestMessage
+    ) -> tuple[Messages, float, bool, bool]:
+        await self.exec_tool_calls(action, exec_timeout=0.0001)
+
+        return [], 0.0, False, False
+
+
+@pytest.mark.asyncio
+async def test_tool_exec_timeout() -> None:
+    env = SlowEnv()
+    _, tools = await env.reset()
+
+    for tool in tools:
+        action = ToolRequestMessage(tool_calls=[ToolCall.from_tool(tool)])
+        with pytest.raises(asyncio.TimeoutError):
+            await env.step(action)
+
+
 class TestRendering:
     class SomeState(BaseModel):
         field: int