Log tool execution start/end time (#304)

Author: sidnarayanan

src/aviary/env.py

@@ -7,6 +7,7 @@
 import logging
 import os
 import random
+import time
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Iterator
 from copy import deepcopy
@@ -216,6 +217,7 @@ async def exec_tool_calls(
         """
 
         async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage:
+            start = time.monotonic()
             try:
                 tool = next(
                     t for t in self.tools if t.info.name == tool_call.function.name
@@ -283,7 +285,14 @@ async def _exec_tool_call(tool_call: ToolCall) -> ToolResponseMessage:
                 s_content = content.model_dump_json(exclude_none=True, by_alias=True)
             else:  # Fallback when content is another type, or None
                 s_content = json.dumps(content)
-            return ToolResponseMessage.from_call(tool_call, content=s_content)
+            return ToolResponseMessage.from_call(
+                tool_call,
+                content=s_content,
+                info={
+                    "start_ts": start,
+                    "end_ts": time.monotonic(),
+                },
+            )
 
         invalid_responses = []
         valid_action = message

src/aviary/tools/base.py

@@ -165,8 +165,10 @@ class ToolResponseMessage(Message):
     )
 
     @classmethod
-    def from_call(cls, call: ToolCall, content: str) -> Self:
-        return cls(content=content, name=call.function.name, tool_call_id=call.id)
+    def from_call(cls, call: ToolCall, content: str, **kwargs) -> Self:
+        return cls(
+            content=content, name=call.function.name, tool_call_id=call.id, **kwargs
+        )
 
     @classmethod
     def from_request(

tests/test_tools.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import os
 import pickle
@@ -689,6 +690,30 @@ 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_timing(self) -> None:
+        sleep_time = 0.1
+
+        async def sleep_tool_fn() -> None:
+            """Zzz."""
+            await asyncio.sleep(sleep_time)
+
+        tool = Tool.from_function(sleep_tool_fn)
+        tool_calls = [ToolCall.from_tool(tool) for _ in range(3)]
+
+        dummy_env = DummyEnv()
+        dummy_env.tools = [tool]
+
+        responses = await dummy_env.exec_tool_calls(
+            ToolRequestMessage(tool_calls=tool_calls)
+        )
+
+        for resp in responses:
+            assert resp.info, "Expected timing info to be present"
+            assert resp.info["end_ts"] - resp.info["start_ts"] >= sleep_time, (
+                "Expected non-trivial time elapsed."
+            )
+
 
 def test_argref_by_name_basic_usage() -> None:
     class MyState: