inspect-ai 0.3.69__py3-none-any.whl → 0.3.71__py3-none-any.whl

inspect_ai/tool/_tools/_web_browser/_web_browser.py

@@ -3,8 +3,9 @@ from textwrap import dedent
 
 from pydantic import Field
 
+from inspect_ai._util.content import ContentText
 from inspect_ai._util.error import PrerequisiteError
-from inspect_ai.tool._tool import Tool, ToolError, tool
+from inspect_ai.tool._tool import Tool, ToolError, ToolResult, tool
 from inspect_ai.tool._tool_call import ToolCall, ToolCallContent, ToolCallView
 from inspect_ai.tool._tool_info import parse_tool_info
 from inspect_ai.tool._tool_with import tool_with
@@ -58,10 +59,10 @@ def web_browser_go() -> Tool:
        Web browser navigation tool.
     """
 
-    async def execute(url: str) -> str:
+    async def execute(url: str) -> ToolResult:
         """Navigate the web browser to a URL.
 
-        Once you have navigated to a page, you will be presented with a web accessibilty tree of the elements on the page. Each element has an ID, which is displayed in brackets at the beginning of its line. For example:
+        Once you have navigated to a page, you will be presented with a web accessibility tree of the elements on the page. Each element has an ID, which is displayed in brackets at the beginning of its line. For example:
 
         ```
         [1] RootWebArea "Google" [focused: True, url: https://www.google.com/]
@@ -99,16 +100,17 @@ def go_without_interactive_docs(tool: Tool) -> Tool:
 
 
 # custom viewer for interactive tool calls that shows a truncated
-# version of current the web accessiblity tree if available
+# version of current the web accessibility tree if available
 
 
 class WebBrowserStore(StoreModel):
+    main_content: str = Field(default_factory=str)
     web_at: str = Field(default_factory=str)
     session_id: str = Field(default_factory=str)
 
 
 def web_at_viewer(call: ToolCall) -> ToolCallView:
-    # get the web accessiblity tree, if we have it create a view from it
+    # get the web accessibility tree, if we have it create a view from it
     web_at = store_as(WebBrowserStore).web_at
     element_id = call.arguments.get("element_id", 0)
     if web_at and element_id:
@@ -141,10 +143,10 @@ def web_browser_click() -> Tool:
        Web browser clicking tool.
     """
 
-    async def execute(element_id: int) -> str:
+    async def execute(element_id: int) -> ToolResult:
         """Click an element on the page currently displayed by the web browser.
 
-        For example, with the following web accessibilty tree:
+        For example, with the following web accessibility tree:
 
         ```
         [304] RootWebArea "Poetry Foundation" [focused: True, url: https://www.poetryfoundation.org/]
@@ -176,7 +178,7 @@ def web_browser_type_submit() -> Tool:
        Web browser type and submit tool.
     """
 
-    async def execute(element_id: int, text: str) -> str:
+    async def execute(element_id: int, text: str) -> ToolResult:
         """Type text into a form input on a web browser page and press ENTER to submit the form.
 
         For example, to execute a search for "Yeats" from this page:
@@ -214,7 +216,7 @@ def web_browser_type() -> Tool:
        Web browser typing tool.
     """
 
-    async def execute(element_id: int, text: str) -> str:
+    async def execute(element_id: int, text: str) -> ToolResult:
         """Type text into an input on a web browser page.
 
         For example, to type "Norah" into the "First Name" search box on this page:
@@ -252,7 +254,7 @@ def web_browser_scroll() -> Tool:
        Web browser scrolling tool.
     """
 
-    async def execute(direction: str) -> str:
+    async def execute(direction: str) -> ToolResult:
         """Scroll the web browser up or down by one page.
 
         Occasionally some very long pages don't display all of their content at once. To see additional content you can scroll the page down with:
@@ -282,7 +284,7 @@ def web_browser_back() -> Tool:
        Web browser back navigation tool.
     """
 
-    async def execute() -> str:
+    async def execute() -> ToolResult:
         """Navigate the web browser back in the browser history.
 
         If you want to view a page that you have previously browsed (or perhaps just didn't find what you were looking for on a page and want to backtrack) use the web_browser_back tool.
@@ -303,7 +305,7 @@ def web_browser_forward() -> Tool:
        Web browser forward navigation tool.
     """
 
-    async def execute() -> str:
+    async def execute() -> ToolResult:
         """Navigate the web browser forward in the browser history.
 
         If you have navigated back in the browser history and then want to navigate forward use the web_browser_forward tool.
@@ -324,7 +326,7 @@ def web_browser_refresh() -> Tool:
        Web browser page refresh tool.
     """
 
-    async def execute() -> str:
+    async def execute() -> ToolResult:
         """Refresh the current page of the web browser.
 
         If you have interacted with a page by clicking buttons and want to reset it to its original state, use the web_browser_refresh tool.
@@ -341,7 +343,7 @@ WEB_CLIENT_REQUEST = "/app/web_browser/web_client.py"
 WEB_CLIENT_NEW_SESSION = "/app/web_browser/web_client_new_session.py"
 
 
-async def web_browser_cmd(cmd: str, *args: str) -> str:
+async def web_browser_cmd(cmd: str, *args: str) -> ToolResult:
     sandbox_env = await sandbox_with(WEB_CLIENT_NEW_SESSION)
     session_flag = ""
     if sandbox_env:
@@ -379,17 +381,30 @@ async def web_browser_cmd(cmd: str, *args: str) -> str:
         if "error" in response and response.get("error", "").strip() != "":
             raise ToolError(str(response.get("error")) or "(unknown error)")
         elif "web_at" in response:
+            main_content = str(response.get("main_content")) or None
             web_at = (
-                str(response.get("web_at")) or "(no web accessiblity tree available)"
+                str(response.get("web_at")) or "(no web accessibility tree available)"
             )
             # Remove base64 data from images.
             web_at_lines = web_at.split("\n")
             web_at_lines = [
                 line.partition("data:image/png;base64")[0] for line in web_at_lines
             ]
-            web_at = "\n".join(web_at_lines)
+
+            store_as(WebBrowserStore).main_content = (
+                main_content or "(no main text summary)"
+            )
             store_as(WebBrowserStore).web_at = web_at
-            return web_at
+
+            web_at = "\n".join(web_at_lines)
+            return (
+                [
+                    ContentText(text=f"main content:\n{main_content}\n\n"),
+                    ContentText(text=f"accessibility tree:\n{web_at}"),
+                ]
+                if main_content
+                else web_at
+            )
         else:
             raise RuntimeError(
                 f"web_browser output must contain either 'error' or 'web_at' field: {result.stdout}"
@@ -425,7 +440,9 @@ async def web_browser_sandbox() -> SandboxEnvironment:
 
 
 def parse_web_browser_output(output: str) -> dict[str, str]:
-    response: dict[str, str] = dict(web_url="", web_at="", info="", error="")
+    response: dict[str, str] = dict(
+        web_url="", main_content="", web_at="", info="", error=""
+    )
     active_field: str | None = None
     active_field_lines: list[str] = []
 
@@ -435,7 +452,9 @@ def parse_web_browser_output(output: str) -> dict[str, str]:
         active_field_lines.clear()
 
     for line in output.splitlines():
-        field_match = re.match(r"^(error|web_at|web_url|info)\s*:\s*(.+)$", line)
+        field_match = re.match(
+            r"^(error|main_content|web_at|web_url|info)\s*:\s*(.+)$", line
+        )
         if field_match:
             collect_active_field()
             active_field = field_match.group(1)

inspect_ai/tool/_tools/_web_search.py

@@ -1,6 +1,6 @@
 import asyncio
 import os
-from typing import Literal, Protocol, cast, runtime_checkable
+from typing import Literal, Protocol, runtime_checkable
 
 import httpx
 from bs4 import BeautifulSoup, NavigableString
@@ -90,8 +90,8 @@ def web_search(
                 return_exceptions=True,
             )
             for page, link in zip(pages, links):
-                if page and not isinstance(page, Exception):
-                    page_contents.append(cast(str, page))
+                if page and not isinstance(page, BaseException):
+                    page_contents.append(page)
                     urls.append(link.url)
                     snippets.append(link.snippet)
             search_calls += 1

inspect_ai/util/__init__.py

@@ -2,7 +2,7 @@ from inspect_ai._util.trace import trace_action, trace_message
 
 from ._concurrency import concurrency
 from ._console import input_screen
-from ._display import DisplayType, display_type
+from ._display import DisplayType, display_counter, display_type
 from ._panel import InputPanel, input_panel
 from ._resource import resource
 from ._sandbox import (
@@ -31,6 +31,7 @@ __all__ = [
     "ExecResult",
     "concurrency",
     "DisplayType",
+    "display_counter",
     "display_type",
     "InputPanel",
     "input_panel",

inspect_ai/util/_concurrency.py

@@ -1,13 +1,19 @@
 import asyncio
+import contextlib
+import time
 from dataclasses import dataclass
+from typing import AsyncIterator
 
+from inspect_ai._util.working import report_sample_waiting_time
 
-def concurrency(
+
+@contextlib.asynccontextmanager
+async def concurrency(
     name: str,
     concurrency: int,
     key: str | None = None,
-) -> asyncio.Semaphore:
-    """Obtain a concurrency context.
+) -> AsyncIterator[None]:
+    """Concurrency context manager.
 
     A concurrency context can be used to limit the number of coroutines
     executing a block of code (e.g calling an API). For example, here
@@ -32,9 +38,6 @@ def concurrency(
          Used if the unique key isn't human readable -- e.g. includes
          api tokens or account ids so that the more readable `name`
          can be presented to users e.g in console UI>
-
-    Returns:
-       Asyncio Semaphore for concurrency context.
     """
     # sort out key
     key = key if key else name
@@ -47,8 +50,11 @@ def concurrency(
         )
         _concurrency_semaphores[key] = semaphore
 
-    # return the semaphore
-    return semaphore.semaphore
+    # wait and yield to protected code
+    start_wait = time.monotonic()
+    async with semaphore.semaphore:
+        report_sample_waiting_time(time.monotonic() - start_wait)
+        yield
 
 
 def concurrency_status() -> dict[str, tuple[int, int]]:

inspect_ai/util/_display.py

@@ -54,3 +54,15 @@ def display_type() -> DisplayType:
 def display_type_initialized() -> bool:
     global _display_type
     return _display_type is not None
+
+
+def display_counter(caption: str, value: str) -> None:
+    """Display a counter in the UI.
+
+    Args:
+        caption: The counter's caption e.g. "HTTP rate limits".
+        value: The counter's value e.g. "42".
+    """
+    from inspect_ai._display.core.active import display
+
+    display().display_counter(caption, value)

inspect_ai/util/_sandbox/context.py

@@ -5,6 +5,7 @@ from typing import Any, NoReturn, cast
 from shortuuid import uuid
 
 from inspect_ai._util.constants import SANDBOX_SETUP_TIMEOUT
+from inspect_ai.util._sandbox.events import SandboxEnvironmentProxy
 
 from .environment import (
     SampleCleanup,
@@ -132,6 +133,9 @@ async def init_sandbox_environments_sample(
     # verify that there is at least one environment and a 'default' env
     validate_sandbox_environments(sandboxenv_type, environments)
 
+    # proxy environments (for recording SandboxEvent)
+    environments = {k: SandboxEnvironmentProxy(v) for k, v in environments.items()}
+
     try:
         # copy files into environments
         await copy_sandbox_environment_files(files, environments)
@@ -148,6 +152,7 @@ async def init_sandbox_environments_sample(
         return environments
 
     except Exception as ex:
+        environments = unproxy_environments(environments)
         await sample_cleanup(task_name, config, environments, True)
         raise ex
 
@@ -161,9 +166,19 @@ async def cleanup_sandbox_environments_sample(
 ) -> None:
     sandboxenv_type = registry_find_sandboxenv(type)
     sample_cleanup = cast(SampleCleanup, getattr(sandboxenv_type, "sample_cleanup"))
+    environments = unproxy_environments(environments)
     await sample_cleanup(task_name, config, environments, interrupted)
 
 
+def unproxy_environments(
+    environments: dict[str, SandboxEnvironment],
+) -> dict[str, SandboxEnvironment]:
+    return {
+        k: v._sandbox
+        for k, v in cast(dict[str, SandboxEnvironmentProxy], environments).items()
+    }
+
+
 async def copy_sandbox_environment_files(
     files: dict[str, bytes], environments: dict[str, SandboxEnvironment]
 ) -> None:

inspect_ai/util/_sandbox/docker/docker.py

@@ -5,7 +5,7 @@ import os
 import tempfile
 from logging import getLogger
 from pathlib import Path, PurePosixPath
-from typing import Literal, Union, cast, overload
+from typing import Literal, Union, overload
 
 from typing_extensions import override
 
@@ -145,7 +145,7 @@ class DockerSandboxEnvironment(SandboxEnvironment):
         project = await ComposeProject.create(
             name=task_project_name(task_name),
             config=config,
-            sample_id=sample.id if sample is not None else None,
+            sample_id=sample.sample.id if sample is not None else None,
             epoch=sample.epoch if sample is not None else None,
             env=env,
         )
@@ -221,9 +221,11 @@ class DockerSandboxEnvironment(SandboxEnvironment):
         # (this enables us to show output for the cleanup operation)
         if not interrupted:
             # extract project from first environment
-            project = cast(
-                DockerSandboxEnvironment, next(iter(environments.values()))
-            )._project
+            project = (
+                next(iter(environments.values()))
+                .as_type(DockerSandboxEnvironment)
+                ._project
+            )
             # cleanup the project
             await project_cleanup(project=project, quiet=True)
 

inspect_ai/util/_sandbox/environment.py

@@ -2,12 +2,24 @@ from __future__ import annotations
 
 import abc
 from dataclasses import dataclass, field
-from typing import Any, Awaitable, Callable, Literal, NamedTuple, Union, overload
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Literal,
+    NamedTuple,
+    Type,
+    TypeVar,
+    Union,
+    overload,
+)
 
 from pydantic import BaseModel, Field
 
 from .._subprocess import ExecResult
 
+ST = TypeVar("ST", bound="SandboxEnvironment")
+
 TaskInit = Callable[[str, Union["SandboxEnvironmentConfigType", None]], Awaitable[None]]
 TaskCleanup = Callable[
     [str, Union["SandboxEnvironmentConfigType", None], bool], Awaitable[None]
@@ -180,6 +192,25 @@ class SandboxEnvironment(abc.ABC):
         """
         raise NotImplementedError("connection not implemented")
 
+    def as_type(self, sandbox_cls: Type[ST]) -> ST:
+        """Verify and return a reference to a subclass of SandboxEnvironment.
+
+        Args:
+           sandbox_cls: Class of sandbox (subclass of SandboxEnvironment)
+
+        Returns:
+           Reference to the sandbox using the requested type.
+
+        Raises:
+           TypeError: If the sandbox is not of the requested type.
+        """
+        if isinstance(self, sandbox_cls):
+            return self
+        else:
+            raise TypeError(
+                f"Expected instance of {sandbox_cls.__name__}, got {type(self).__name__}"
+            )
+
     @classmethod
     def config_files(cls) -> list[str]:
         """Standard config files for this provider (used for automatic discovery)"""

inspect_ai/util/_sandbox/events.py

@@ -0,0 +1,183 @@
+import contextlib
+import shlex
+from datetime import datetime
+from typing import Iterator, Literal, Type, Union, overload
+
+from pydantic import JsonValue
+from pydantic_core import to_jsonable_python
+from typing_extensions import override
+
+from inspect_ai._util.text import truncate_lines
+from inspect_ai.util._subprocess import ExecResult
+
+from .environment import (
+    ST,
+    SandboxConnection,
+    SandboxEnvironment,
+    SandboxEnvironmentConfigType,
+)
+
+
+class SandboxEnvironmentProxy(SandboxEnvironment):
+    def __init__(self, sandbox: SandboxEnvironment) -> None:
+        self._sandbox = sandbox
+        self._events = True
+
+    @override
+    async def exec(
+        self,
+        cmd: list[str],
+        input: str | bytes | None = None,
+        cwd: str | None = None,
+        env: dict[str, str] = {},
+        user: str | None = None,
+        timeout: int | None = None,
+        timeout_retry: bool = True,
+    ) -> ExecResult[str]:
+        from inspect_ai.log._transcript import SandboxEvent, transcript
+
+        # started
+        timestamp = datetime.now()
+
+        # make call
+        result = await self._sandbox.exec(
+            cmd, input, cwd, env, user, timeout, timeout_retry
+        )
+
+        # yield event
+        options: dict[str, JsonValue] = {}
+        if cwd:
+            options["cwd"] = cwd
+        if env:
+            options["env"] = to_jsonable_python(env)
+        if user:
+            options["user"] = user
+        if timeout is not None:
+            options["timeout"] = timeout
+        if timeout_retry is not True:
+            options["timeout_retry"] = timeout_retry
+
+        if self._events:
+            transcript()._event(
+                SandboxEvent(
+                    timestamp=timestamp,
+                    action="exec",
+                    cmd=" ".join([shlex.quote(c) for c in cmd]),
+                    input=content_display(input) if input is not None else None,
+                    options=options,
+                    result=result.returncode,
+                    output=content_display(
+                        f"{result.stderr}\n\n{result.stdout}"
+                        if result.stderr
+                        else result.stdout
+                    ),
+                    completed=datetime.now(),
+                )
+            )
+
+        # return result
+        return result
+
+    @override
+    async def write_file(self, file: str, contents: str | bytes) -> None:
+        from inspect_ai.log._transcript import SandboxEvent, transcript
+
+        timestamp = datetime.now()
+
+        # make call
+        await self._sandbox.write_file(file, contents)
+
+        # yield event
+        if self._events:
+            transcript()._event(
+                SandboxEvent(
+                    timestamp=timestamp,
+                    action="write_file",
+                    file=file,
+                    input=content_display(contents),
+                    completed=datetime.now(),
+                )
+            )
+
+    @overload
+    async def read_file(self, file: str, text: Literal[True] = True) -> str: ...
+
+    @overload
+    async def read_file(self, file: str, text: Literal[False]) -> bytes: ...
+
+    @override
+    async def read_file(self, file: str, text: bool = True) -> Union[str | bytes]:
+        from inspect_ai.log._transcript import SandboxEvent, transcript
+
+        timestamp = datetime.now()
+
+        # make call
+        if text is True:
+            output: str | bytes = await self._sandbox.read_file(file, True)
+        else:
+            output = await self._sandbox.read_file(file, False)
+
+        # yield event
+        if self._events:
+            transcript()._event(
+                SandboxEvent(
+                    timestamp=timestamp,
+                    action="read_file",
+                    file=file,
+                    output=content_display(output),
+                    completed=datetime.now(),
+                )
+            )
+
+        # return result
+        return output
+
+    @override
+    async def connection(self) -> SandboxConnection:
+        return await self._sandbox.connection()
+
+    @override
+    def as_type(self, sandbox_cls: Type[ST]) -> ST:
+        if isinstance(self._sandbox, sandbox_cls):
+            return self._sandbox
+        else:
+            raise TypeError(
+                f"Expected instance of {sandbox_cls.__name__}, got {type(self._sandbox).__name__}"
+            )
+
+    @contextlib.contextmanager
+    def no_events(self) -> Iterator[None]:
+        self._events = False
+        try:
+            yield
+        finally:
+            self._events = True
+
+    @classmethod
+    async def sample_cleanup(
+        cls,
+        task_name: str,
+        config: SandboxEnvironmentConfigType | None,
+        environments: dict[str, SandboxEnvironment],
+        interrupted: bool,
+    ) -> None:
+        pass
+
+
+def content_display(content: str | bytes) -> str:
+    if isinstance(content, str):
+        content, truncated = truncate_lines(content, 20)
+        if truncated:
+            content = f"{content}\n\nOutput truncated ({truncated} additional lines)"
+        return content
+    else:
+        return f"binary ({pretty_size(len(content))})"
+
+
+def pretty_size(size: int) -> str:
+    if size < 1024:
+        return f"{size} B"
+    if size < 1024 * 1024:
+        return f"{size / 1024:.2f} KB"
+
+    return f"{size / (1024 * 1024):.2f} MB"

inspect_ai/util/_sandbox/local.py

@@ -1,7 +1,7 @@
 import tempfile
 import warnings
 from pathlib import Path
-from typing import Literal, Union, cast, overload
+from typing import Literal, Union, overload
 
 from typing_extensions import override
 
@@ -40,8 +40,8 @@ class LocalSandboxEnvironment(SandboxEnvironment):
         interrupted: bool,
     ) -> None:
         for environment in environments.values():
-            env = cast(LocalSandboxEnvironment, environment)
-            env.directory.cleanup()
+            sandbox = environment.as_type(LocalSandboxEnvironment)
+            sandbox.directory.cleanup()
 
     def __init__(self) -> None:
         self.directory = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)