inspect-ai 0.3.70__py3-none-any.whl → 0.3.72__py3-none-any.whl

inspect_ai/tool/_tools/_web_browser/_resources/web_server.py

@@ -1,52 +1,191 @@
-"""Simple script to run and test the RPC server."""
+import threading
+from typing import Awaitable, Callable, Unpack
 
-from concurrent import futures
+from aiohttp.web import Application, Request, Response, run_app
+from jsonrpcserver import Result, Success, async_dispatch, method
 
-import dm_env_servicer
-import grpc
-import web_environment
-from dm_env_rpc.v1 import connection as dm_env_connection
-from dm_env_rpc.v1 import dm_env_rpc_pb2, dm_env_rpc_pb2_grpc
+from constants import DEFAULT_SESSION_NAME, SERVER_PORT
+from playwright_browser import PlaywrightBrowser
+from playwright_crawler import PlaywrightCrawler
+from scale_factor import get_screen_scale_factor
+from web_browser_rpc_types import (
+    ClickArgs,
+    CrawlerBaseArgs,
+    CrawlerResponse,
+    GoArgs,
+    NewSessionArgs,
+    NewSessionResponse,
+    ScrollArgs,
+    TypeOrSubmitArgs,
+)
 
-_DM_ENV_BASE_PORT = 9443
 
+class Sessions:
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._browser: PlaywrightBrowser | None = None
+        self._sessions: dict[str, PlaywrightCrawler] = {}
 
-def main():
-    # Keep connection open even when empty data pings are received. This is
-    # required to avoid a "too many pings" error.
-    options = [
-        ("grpc.keepalive_permit_without_calls", 0),
-        ("grpc.keepalive_timeout_ms", 20000),
-        ("grpc.http2.max_pings_without_data", 0),
-        ("grpc.http2.max_ping_strikes", 0),
-        ("grpc.http2.min_recv_ping_interval_without_data_ms", 0),
-    ]
-    grpc_server = grpc.server(
-        # We must have a single worker thread since the web environment is not
-        # thread safe.
-        futures.ThreadPoolExecutor(max_workers=1),
-        options=options,
-    )
-    env_service = dm_env_servicer.EnvironmentService(web_environment.WebEnvironment)
-    dm_env_rpc_pb2_grpc.add_EnvironmentServicer_to_server(env_service, grpc_server)
+    async def new_session(self, headful: bool) -> str:
+        with self._lock:
+            if not self._browser:
+                self._browser = await PlaywrightBrowser.create(headless=not headful)
+            current_count = len(self._sessions)
+            name = (
+                DEFAULT_SESSION_NAME
+                if current_count == 0
+                else f"{DEFAULT_SESSION_NAME}_{current_count}"
+            )
+            crawler = await PlaywrightCrawler.create(
+                await self._browser.get_new_context(),
+                device_scale_factor=get_screen_scale_factor() if headful else 1,
+            )
+            self._sessions[name] = crawler
+            return name
 
-    grpc_server.add_secure_port(
-        f"localhost:{_DM_ENV_BASE_PORT}", grpc.local_server_credentials()
-    )
+    async def get_crawler_for_session(self, name: str) -> PlaywrightCrawler:
+        if not self._sessions:
+            await self.new_session(False)
+        return self._sessions[name]
 
-    grpc_server.start()
 
-    # Creating a world with the headless browser after the server started.
-    channel = grpc.secure_channel(
-        f"localhost:{_DM_ENV_BASE_PORT}",
-        grpc.local_channel_credentials(),
-        options=options,
+sessions = Sessions()
+
+
+@method
+async def new_session(**kwargs: Unpack[NewSessionArgs]) -> NewSessionResponse:
+    return Success(
+        NewSessionResponse(
+            session_name=await sessions.new_session(kwargs.get("headful", False))
+        ).model_dump()
     )
-    connection = dm_env_connection.Connection(channel)
-    connection.send(dm_env_rpc_pb2.CreateWorldRequest())
-    connection.close()
 
-    grpc_server.wait_for_termination()
+
+@method
+async def web_go(**kwargs: Unpack[GoArgs]) -> Result:
+    async def handler(crawler: PlaywrightCrawler):
+        await (await crawler.current_page).go_to_url(kwargs["url"])
+
+    return await _execute_crawler_command(kwargs["session_name"], handler)
+
+
+@method
+async def web_click(**kwargs: Unpack[ClickArgs]) -> Result:
+    async def handler(crawler: PlaywrightCrawler):
+        await (await crawler.current_page).click(kwargs["element_id"])
+
+    return await _execute_crawler_command(kwargs["session_name"], handler)
+
+
+@method
+async def web_scroll(**kwargs: Unpack[ScrollArgs]) -> Result:
+    async def handler(crawler: PlaywrightCrawler):
+        await (await crawler.current_page).scroll(kwargs["direction"])
+
+    return await _execute_crawler_command(kwargs["session_name"], handler)
+
+
+@method
+async def web_forward(**kwargs: Unpack[CrawlerBaseArgs]) -> Result:
+    async def handler(crawler: PlaywrightCrawler):
+        await (await crawler.current_page).forward()
+
+    return await _execute_crawler_command(kwargs["session_name"], handler)
+
+
+@method
+async def web_back(**kwargs: Unpack[CrawlerBaseArgs]) -> Result:
+    async def handler(crawler: PlaywrightCrawler):
+        await (await crawler.current_page).back()
+
+    return await _execute_crawler_command(kwargs["session_name"], handler)
+
+
+@method
+async def web_refresh(**kwargs: Unpack[CrawlerBaseArgs]) -> Result:
+    async def handler(crawler: PlaywrightCrawler):
+        await (await crawler.current_page).refresh()
+
+    return await _execute_crawler_command(kwargs["session_name"], handler)
+
+
+@method
+async def web_type(**kwargs: Unpack[TypeOrSubmitArgs]) -> Result:
+    async def handler(crawler: PlaywrightCrawler):
+        await (await crawler.current_page).type(
+            kwargs["element_id"], _str_from_str_or_list(kwargs["text"])
+        )
+
+    return await _execute_crawler_command(kwargs["session_name"], handler)
+
+
+@method
+async def web_type_submit(**kwargs: Unpack[TypeOrSubmitArgs]) -> Result:
+    async def handler(crawler: PlaywrightCrawler):
+        await (await crawler.current_page).clear(kwargs["element_id"])
+        await (await crawler.current_page).type(
+            kwargs["element_id"], _str_from_str_or_list(kwargs["text"]) + "\n"
+        )
+
+    return await _execute_crawler_command(kwargs["session_name"], handler)
+
+
+async def _execute_crawler_command(
+    session_name: str, handler: Callable[[PlaywrightCrawler], Awaitable[None]]
+) -> Result:
+    if not sessions:
+        await new_session()
+    try:
+        crawler = await sessions.get_crawler_for_session(session_name)
+        await handler(crawler)
+        await (await crawler.current_page).update()
+
+        # If there's a cookies message click to sort it out.
+        await _auto_click_cookies(crawler)
+
+        return Success(
+            CrawlerResponse(
+                web_url=(await crawler.current_page).url.split("?")[0],
+                main_content=(await crawler.current_page).render_main_content(),
+                web_at=(await crawler.current_page).render_at(),
+                error=None,
+            ).model_dump()
+        )
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        return Success(
+            CrawlerResponse(
+                web_url=(await crawler.current_page).url.split("?")[0],
+                web_at="encountered error",
+                error=str(e),
+            ).model_dump()
+        )
+
+
+def _str_from_str_or_list(str_or_list: str | list[str]) -> str:
+    return str_or_list if isinstance(str_or_list, str) else " ".join(str_or_list)
+
+
+async def _auto_click_cookies(crawler: PlaywrightCrawler):
+    """Autoclick any cookies popup."""
+    try:
+        accept_node = (await crawler.current_page).lookup_node("<Accept all>")
+    except LookupError:
+        return
+    await (await crawler.current_page).click(accept_node.node_id)
+    await (await crawler.current_page).update()
+
+
+def main():
+    async def handle_request(request: Request) -> Response:
+        return Response(
+            text=await async_dispatch(await request.text()),
+            content_type="application/json",
+        )
+
+    app = Application()
+    app.router.add_post("/", handle_request)
+
+    run_app(app, port=SERVER_PORT)
 
 
 if __name__ == "__main__":

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/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/_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/events.py

@@ -1,5 +1,7 @@
+import contextlib
 import shlex
-from typing import Literal, Type, Union, overload
+from datetime import datetime
+from typing import Iterator, Literal, Type, Union, overload
 
 from pydantic import JsonValue
 from pydantic_core import to_jsonable_python
@@ -19,6 +21,7 @@ from .environment import (
 class SandboxEnvironmentProxy(SandboxEnvironment):
     def __init__(self, sandbox: SandboxEnvironment) -> None:
         self._sandbox = sandbox
+        self._events = True
 
     @override
     async def exec(
@@ -33,6 +36,9 @@ class SandboxEnvironmentProxy(SandboxEnvironment):
     ) -> 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
@@ -50,20 +56,24 @@ class SandboxEnvironmentProxy(SandboxEnvironment):
             options["timeout"] = timeout
         if timeout_retry is not True:
             options["timeout_retry"] = timeout_retry
-        transcript()._event(
-            SandboxEvent(
-                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
-                ),
+
+        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
@@ -72,15 +82,22 @@ class SandboxEnvironmentProxy(SandboxEnvironment):
     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
-        transcript()._event(
-            SandboxEvent(
-                action="write_file", file=file, input=content_display(contents)
+        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: ...
@@ -92,6 +109,8 @@ class SandboxEnvironmentProxy(SandboxEnvironment):
     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)
@@ -99,9 +118,16 @@ class SandboxEnvironmentProxy(SandboxEnvironment):
             output = await self._sandbox.read_file(file, False)
 
         # yield event
-        transcript()._event(
-            SandboxEvent(action="read_file", file=file, output=content_display(output))
-        )
+        if self._events:
+            transcript()._event(
+                SandboxEvent(
+                    timestamp=timestamp,
+                    action="read_file",
+                    file=file,
+                    output=content_display(output),
+                    completed=datetime.now(),
+                )
+            )
 
         # return result
         return output
@@ -119,6 +145,14 @@ class SandboxEnvironmentProxy(SandboxEnvironment):
                 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,