inspect-ai 0.3.90__py3-none-any.whl → 0.3.91__py3-none-any.whl

inspect_ai/tool/_tool_support_helpers.py

@@ -4,197 +4,173 @@ This module provides helper code for handling JSON-RPC communication between the
 It includes definitions for JSON-RPC request and response models, as well as functions to create and parse JSON-RPC requests and responses.
 """
 
-import json
-from itertools import count
 from textwrap import dedent
-from typing import Literal, Type, TypeVar, cast
-
-from pydantic import BaseModel, RootModel
+from typing import Type
 
 from inspect_ai._util.error import PrerequisiteError
-from inspect_ai.tool._tool import ToolError, ToolParsingError
 from inspect_ai.util import sandbox_with
 from inspect_ai.util._sandbox.environment import SandboxEnvironment
 
+from ._json_rpc_helpers import (
+    BaseModelT,
+    JSONRPCParamsType,
+    JSONRPCTransport,
+    ScalarT,
+    _rpc_call_description,
+    create_json_rpc_request,
+)
+from ._json_rpc_helpers import exec_model_request as model_request
+from ._json_rpc_helpers import exec_notification as notification_helper
+from ._json_rpc_helpers import exec_scalar_request as scalar_request
 
-class JSONRPCResponseBase(BaseModel):
-    jsonrpc: Literal["2.0"]
-    id: int | float | str
-
-
-class JSONRPCSuccessResponse(JSONRPCResponseBase):
-    result: object
-
-
-class JSONRPCError(BaseModel):
-    """See: https://www.jsonrpc.org/specification#error_object"""
-
-    code: int
-    message: str
-    data: object | None = None
-
-
-class JSONRPCErrorResponse(JSONRPCResponseBase):
-    error: JSONRPCError
-
-
-class JSONRPCResponse(RootModel[JSONRPCSuccessResponse | JSONRPCErrorResponse]):
-    pass
 
+async def exec_scalar_request(
+    sandbox: SandboxEnvironment,
+    method: str,
+    params: JSONRPCParamsType,
+    result_type: Type[ScalarT],
+    timeout: int | None = None,
+    user: str | None = None,
+) -> ScalarT:
+    return await scalar_request(
+        method,
+        params,
+        result_type,
+        transport=ToolSupportSandboxTransport(sandbox, timeout, user),
+    )
 
-BaseModelT = TypeVar("BaseModelT", bound=BaseModel)
-StrOrModelT = TypeVar("StrOrModelT", bound=str | BaseModel)
 
-id_generator = count(666)
+async def exec_model_request(
+    sandbox: SandboxEnvironment,
+    method: str,
+    params: JSONRPCParamsType,
+    result_type: Type[BaseModelT],
+    timeout: int | None = None,
+    user: str | None = None,
+) -> BaseModelT:
+    return await model_request(
+        method,
+        params,
+        result_type,
+        transport=ToolSupportSandboxTransport(sandbox, timeout, user),
+    )
 
 
-async def exec_sandbox_rpc(
+async def exec_notification(
     sandbox: SandboxEnvironment,
     method: str,
-    params: dict[str, object] | tuple[object, ...],
-    result_cls: Type[StrOrModelT],
+    params: JSONRPCParamsType,
     timeout: int | None = None,
     user: str | None = None,
-) -> StrOrModelT:
-    """
-    Execute a JSON-RPC command to a sandbox environment.
-
-    Note that the JSON RPC request is sent to the exec'ed program via stdin.
+) -> None:
+    return await notification_helper(
+        method, params, transport=ToolSupportSandboxTransport(sandbox, timeout, user)
+    )
 
-    Args:
-      sandbox (SandboxEnvironment): The sandbox environment to execute the command in.
-      method (str): The JSON-RPC method to call.
-      params (dict[str, object] | tuple[object, ...]): The parameters for the JSON-RPC method.
-      result_cls (Type[BaseModelT]): The class to use for parsing the result.
-      timeout (int | None, optional): The timeout for the execution. Defaults to None.
-      user: Optional username or UID to run the command as.
 
-    Returns:
-      BaseModelT: The parsed result of the JSON-RPC call.
+class ToolSupportSandboxTransport(JSONRPCTransport):
+    """
+    A transport callable that uses a sandbox for RPC communication.
 
-    Raises:
-      RuntimeError: If the sandbox execution fails or if there is an error in the JSON-RPC response.
-      ToolParsingError: If the JSON-RPC response contains a specific error code indicating a parsing error.
+    This class implements the TransportCallable protocol and encapsulates
+    the sandbox, timeout, and user parameters needed for sandbox-based
+    RPC communication.
     """
-    exec_result = await sandbox.exec(
-        [SANDBOX_CLI, "exec"],
-        input=_create_json_rpc_request(method, params),
-        timeout=timeout,
-        user=user,
-    )
 
-    if not exec_result.success:
-        raise RuntimeError(
-            f"Sandbox.exec failure executing {_rpc_call_description(method, params)}: {exec_result.stderr}"
+    def __init__(
+        self,
+        sandbox: SandboxEnvironment,
+        timeout: int | None = None,
+        user: str | None = None,
+    ):
+        """
+        Initialize a new SandboxTransportCallable.
+
+        Args:
+            sandbox (SandboxEnvironment): The sandbox environment to use.
+            timeout (int | None, optional): The timeout for executions. Defaults to None.
+            user (str | None, optional): Username or UID to run commands as. Defaults to None.
+        """
+        self.sandbox = sandbox
+        self.timeout = timeout
+        self.user = user
+
+    async def __call__(
+        self, method: str, params: JSONRPCParamsType, is_notification: bool
+    ) -> str:
+        """
+        Execute an RPC request using the sandbox transport.
+
+        Args:
+            method (str): The JSON-RPC method to call.
+            params (dict[str, object] | tuple[object, ...]): The parameters for the JSON-RPC method.
+            is_notification (bool): Whether this is a notification (no response expected).
+
+        Returns:
+            str: The response from the RPC call.
+
+        Raises:
+            RuntimeError: If the sandbox execution fails.
+        """
+        exec_result = await self.sandbox.exec(
+            [SANDBOX_CLI, "exec"],
+            input=create_json_rpc_request(method, params, is_notification),
+            timeout=self.timeout,
+            user=self.user,
         )
 
-    match _parse_json_rpc_response(exec_result.stdout, result_cls):
-        case JSONRPCError(code=-32601 | -32602, message=message):
-            raise ToolParsingError(message)
-        case JSONRPCError(code=-32000, message=message):
-            raise ToolError(message)
-        case JSONRPCError(code=code, message=message):
-            raise RuntimeError(
-                f"Error executing tool command {_rpc_call_description(method, params)}: {code=} {message}"
-            )
-        # case result_cls() as model: yields a mypy error since it has narrowed model down
-        # to BaseModel and not BaseModelT. ???
-        case model if isinstance(model, result_cls):
-            return model
-        case not_possible:
+        if not exec_result.success:
             raise RuntimeError(
-                f"Error executing tool command {_rpc_call_description(method, params)}: {not_possible}"
+                f"Sandbox.exec failure executing {_rpc_call_description(method, params)}: {exec_result.stderr}"
             )
+        return exec_result.stdout
 
 
 SANDBOX_CLI = "inspect-tool-support"
 INSPECT_TOOL_SUPPORT_IMAGE_DOCKERHUB = "aisiuk/inspect-tool-support"
 
 
-async def tool_container_sandbox(tool_name: str) -> SandboxEnvironment:
-    sb = await sandbox_with(SANDBOX_CLI, True)
-    if sb:
+async def tool_container_sandbox(
+    tool_name: str, *, sandbox_name: str | None = None
+) -> SandboxEnvironment:
+    if sb := await sandbox_with(SANDBOX_CLI, True, name=sandbox_name):
         return sb
-    else:
-        msg = dedent(f"""
-                The {tool_name} service was not found in any of the sandboxes for this sample. Please add the {tool_name} to your configuration.
-
-                For example, the following Docker compose file uses the {INSPECT_TOOL_SUPPORT_IMAGE_DOCKERHUB} reference image as its default sandbox:
-
-                services:
-                  default:
-                    image: "{INSPECT_TOOL_SUPPORT_IMAGE_DOCKERHUB}"
-                    init: true
-
-                Alternatively, you can include the service into your own Dockerfile:
-
-                ENV PATH="$PATH:/opt/inspect_tool_support/bin"
-                RUN python -m venv /opt/inspect_tool_support && \\
-                    /opt/inspect_tool_support/bin/pip install inspect-tool-support && \\
-                    /opt/inspect_tool_support/bin/inspect-tool-support post-install
-                """).strip()
-        raise PrerequisiteError(msg)
-
-
-def _create_json_rpc_request(
-    method: str, params: dict[str, object] | tuple[object, ...]
-) -> str:
-    return json.dumps(
-        {
-            "jsonrpc": "2.0",
-            "method": method,
-            "id": next(id_generator),
-            "params": list(params) if isinstance(params, tuple) else params,
-        }
+
+    # This sort of programmatic sentence building will not cut it if we ever
+    # support other languages.
+    raise PrerequisiteError(
+        dedent(f"""
+            The {tool_name} service was not found in {"any of the sandboxes" if sandbox_name is None else f"the sandbox '{sandbox_name}'"} for this sample. Please add the {tool_name} to your configuration.
+
+            For example, the following Docker compose file uses the {INSPECT_TOOL_SUPPORT_IMAGE_DOCKERHUB} reference image as its default sandbox:
+
+            services:
+              default:
+                image: "{INSPECT_TOOL_SUPPORT_IMAGE_DOCKERHUB}"
+                init: true
+
+            Alternatively, you can include the service into your own Dockerfile:
+
+            ENV PATH="$PATH:/opt/inspect_tool_support/bin"
+            RUN python -m venv /opt/inspect_tool_support && \\
+                /opt/inspect_tool_support/bin/pip install inspect-tool-support && \\
+                /opt/inspect_tool_support/bin/inspect-tool-support post-install
+            """).strip()
     )
 
 
-def _rpc_call_description(
-    method: str, params: dict[str, object] | tuple[object, ...]
-) -> str:
+def create_sandbox_transport(
+    sandbox: SandboxEnvironment, timeout: int | None = None, user: str | None = None
+) -> JSONRPCTransport:
     """
-    Generate a string description of an RPC call.
+    Create a transport callable that uses a sandbox for RPC communication.
 
     Args:
-        method (str): The name of the RPC method.
-        params (dict[str, object] | tuple[object, ...]): The parameters for the RPC method.
+        sandbox (SandboxEnvironment): The sandbox environment to use.
+        timeout (int | None, optional): The timeout for executions. Defaults to None.
+        user (str | None, optional): Username or UID to run commands as. Defaults to None.
 
     Returns:
-        str: A string description of the RPC call.
-
-    Examples:
-        >>> _rpc_call_description("subtract", {"minuend": 42, "subtrahend": 23})
-        'subtract(minuend: 42, subtrahend: 23)'
-
-        >>> _rpc_call_description("subtract", (42, 23))
-        'subtract(42, 23)'
+        TransportCallable: A transport callable that conforms to the TransportCallable protocol.
     """
-    normalized_params = (
-        list(map(str, params))
-        if isinstance(params, tuple)
-        else [f"{k}: {v}" for k, v in params.items()]
-    )
-    return f"{method}({', '.join(normalized_params)})"
-
-
-def _parse_json_rpc_response(
-    response_str: str,
-    result_cls: Type[StrOrModelT],
-) -> StrOrModelT | JSONRPCError:
-    match JSONRPCResponse.model_validate_json(response_str).root:
-        case JSONRPCErrorResponse(error=error):
-            return error
-        case JSONRPCSuccessResponse(result=rpc_result):
-            # TODO: Wow. Is there really no way to convince Python to narrow these types
-            # and avoid the cast's
-            if result_cls is str:
-                if not isinstance(rpc_result, str):
-                    raise ValueError(f"Expected string result, got {type(rpc_result)}")
-                return cast(StrOrModelT, rpc_result)
-            else:
-                return cast(
-                    StrOrModelT,
-                    cast(BaseModel, result_cls).model_validate(rpc_result, strict=True),
-                )
-        case _:
-            raise ValueError(f"Unexpected JSON RPC response: {response_str}")
+    return ToolSupportSandboxTransport(sandbox=sandbox, timeout=timeout, user=user)

inspect_ai/tool/_tools/_bash_session.py

@@ -3,7 +3,7 @@ from shortuuid import uuid
 
 from inspect_ai.tool import ToolResult
 from inspect_ai.tool._tool_support_helpers import (
-    exec_sandbox_rpc,
+    exec_model_request,
     tool_container_sandbox,
 )
 from inspect_ai.util import StoreModel, store_as
@@ -98,11 +98,11 @@ def bash_session(*, timeout: int | None = None, instance: str | None = uuid()) -
 
         if not store.session_id:
             store.session_id = (
-                await exec_sandbox_rpc(
-                    sandbox,
-                    "bash_session_new_session",
-                    {},
-                    NewSessionResult,
+                await exec_model_request(
+                    sandbox=sandbox,
+                    method="bash_session_new_session",
+                    params={},
+                    result_type=NewSessionResult,
                     timeout=timeout,
                 )
             ).session_name
@@ -110,11 +110,11 @@ def bash_session(*, timeout: int | None = None, instance: str | None = uuid()) -
         params["session_name"] = store.session_id
 
         result = (
-            await exec_sandbox_rpc(
-                sandbox,
-                "bash_session",
-                params,
-                BashResult,
+            await exec_model_request(
+                sandbox=sandbox,
+                method="bash_session",
+                params=params,
+                result_type=BashResult,
                 timeout=timeout,
             )
         ).root

inspect_ai/tool/_tools/_text_editor.py

@@ -5,7 +5,7 @@ from pydantic import BaseModel, Discriminator, RootModel
 
 from inspect_ai.tool import ToolResult
 from inspect_ai.tool._tool_support_helpers import (
-    exec_sandbox_rpc,
+    exec_scalar_request,
     tool_container_sandbox,
 )
 
@@ -110,11 +110,11 @@ def text_editor(timeout: int | None = None, user: str | None = None) -> Tool:
             if k in inspect.signature(execute).parameters
         }
 
-        return await exec_sandbox_rpc(
-            sandbox,
-            "text_editor",
-            params,
-            TextEditorResult,
+        return await exec_scalar_request(
+            sandbox=sandbox,
+            method="text_editor",
+            params=params,
+            result_type=TextEditorResult,
             timeout=timeout,
         )
 

inspect_ai/tool/_tools/_web_browser/_web_browser.py

@@ -9,7 +9,7 @@ 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_support_helpers import (
-    exec_sandbox_rpc,
+    exec_model_request,
     tool_container_sandbox,
 )
 from inspect_ai.tool._tool_with import tool_with
@@ -414,18 +414,18 @@ async def _web_browser_cmd(
 
     if not store.session_id:
         store.session_id = (
-            await exec_sandbox_rpc(
-                sandbox_env,
-                "web_new_session",
-                {"headful": False},
-                NewSessionResult,
+            await exec_model_request(
+                sandbox=sandbox_env,
+                method="web_new_session",
+                params={"headful": False},
+                result_type=NewSessionResult,
             )
         ).session_name
 
     params["session_name"] = store.session_id
 
-    crawler_result = await exec_sandbox_rpc(
-        sandbox_env, tool_name, params, CrawlerResult
+    crawler_result = await exec_model_request(
+        sandbox=sandbox_env, method=tool_name, params=params, result_type=CrawlerResult
     )
     if crawler_result.error and crawler_result.error.strip() != "":
         raise ToolError(crawler_result.error)

inspect_ai/util/_anyio.py

@@ -1,3 +1,4 @@
+import itertools
 import sys
 
 if sys.version_info < (3, 11):
@@ -5,23 +6,33 @@ if sys.version_info < (3, 11):
 
 
 def inner_exception(exc: Exception) -> Exception:
-    flattended = flatten_exception_group(exc)
-    return flattended[0]
-
-
-def flatten_exception_group(exc: Exception) -> list[Exception]:
-    """Recursively flatten an ExceptionGroup to get all contained exceptions."""
-    if (
-        hasattr(exc, "__context__")
-        and exc.__context__ is not None
-        and isinstance(exc.__context__, Exception)
-    ):
-        return flatten_exception_group(exc.__context__) + [exc]
-
-    if isinstance(exc, ExceptionGroup):
-        flattened = []
-        for nested_exc in exc.exceptions:
-            flattened.extend(flatten_exception_group(nested_exc))
-        return flattened
-
-    return [exc]
+    return _flatten_exception(exc)[0]
+
+
+def _flatten_exception(exc: Exception) -> list[Exception]:
+    """Recursively flatten an exception to get all related (__context__) and contained (ExceptionGroup) exceptions."""
+    context_to_follow = (
+        [exc.__context__]
+        # conceptually, if __cause__ is present, it means that this exception
+        # wraps the cause - rather than cause being a separate error. We'll
+        # follow __context__ only if __cause__ is None
+        if exc.__cause__ is None and isinstance(exc.__context__, Exception)
+        else []
+    )
+
+    (maybe_this_exception, children_to_follow) = (
+        ([], exc.exceptions)
+        # if it's a group, follow the children discarding the group
+        if isinstance(exc, ExceptionGroup)
+        else ([exc], [])
+    )
+
+    # We have to use a set since the same exception is likely to be included in
+    # both __context__ and .exceptions
+    other_exceptions = [
+        flattened_e
+        for e in set(itertools.chain(context_to_follow, children_to_follow))
+        for flattened_e in _flatten_exception(e)
+    ]
+
+    return maybe_this_exception + other_exceptions

inspect_ai/util/_json.py

@@ -2,12 +2,14 @@ import types
 import typing
 from copy import deepcopy
 from dataclasses import is_dataclass
+from datetime import date, datetime, time
 from typing import (
     Any,
     Dict,
     List,
     Literal,
     Optional,
+    Set,
     Tuple,
     Type,
     Union,
@@ -30,6 +32,9 @@ class JSONSchema(BaseModel):
     type: JSONType | None = Field(default=None)
     """JSON type of tool parameter."""
 
+    format: str | None = Field(default=None)
+    """Format of the parameter (e.g. date-time)."""
+
     description: str | None = Field(default=None)
     """Parameter description."""
 
@@ -80,7 +85,13 @@ def json_schema(t: Type[Any]) -> JSONSchema:
             return JSONSchema(type="string")
         elif t is bool:
             return JSONSchema(type="boolean")
-        elif t is list:
+        elif t is datetime:
+            return JSONSchema(type="string", format="date-time")
+        elif t is date:
+            return JSONSchema(type="string", format="date")
+        elif t is time:
+            return JSONSchema(type="string", format="time")
+        elif t is list or t is set:
             return JSONSchema(type="array", items=JSONSchema())
         elif t is dict:
             return JSONSchema(type="object", additionalProperties=JSONSchema())
@@ -94,7 +105,14 @@ def json_schema(t: Type[Any]) -> JSONSchema:
             return JSONSchema(type="null")
         else:
             return JSONSchema()
-    elif origin is list or origin is List or origin is tuple or origin is Tuple:
+    elif (
+        origin is list
+        or origin is List
+        or origin is tuple
+        or origin is Tuple
+        or origin is set
+        or origin is Set
+    ):
         return JSONSchema(
             type="array", items=json_schema(args[0]) if args else JSONSchema()
         )

inspect_ai/util/_sandbox/context.py

@@ -24,7 +24,7 @@ def sandbox(name: str | None = None) -> SandboxEnvironment:
     """Get the SandboxEnvironment for the current sample.
 
     Args:
-      name (str | None): Optional sandbox environmnent name.
+      name (str | None): Optional sandbox environment name.
 
     Return:
       SandboxEnvironment instance.
@@ -45,12 +45,14 @@ def sandbox(name: str | None = None) -> SandboxEnvironment:
         environment = environments.get(name, None)
         if not environment:
             raise ValueError(
-                f"SandboxEnvironment '{name}' is not a recoginized environment name."
+                f"SandboxEnvironment '{name}' is not a recognized environment name."
             )
         return environment
 
 
-async def sandbox_with(file: str, on_path: bool = False) -> SandboxEnvironment | None:
+async def sandbox_with(
+    file: str, on_path: bool = False, *, name: str | None = None
+) -> SandboxEnvironment | None:
     """Get the SandboxEnvironment for the current sample that has the specified file.
 
     Args:
@@ -58,9 +60,12 @@ async def sandbox_with(file: str, on_path: bool = False) -> SandboxEnvironment |
         True, file should be a filename that exists on the system path.
       on_path (bool): If True, file is a filename to be verified using "which".
         If False, file is a path to be checked within the sandbox environments.
+      name (str | None): Optional sandbox environment name.
+
 
     Return:
-      SandboxEnvironment instance or None if no sandboxes had the file.
+      SandboxEnvironment instance or None if none of the sandboxes (or the named
+      sandbox) had the file.
     """
     # get environments and with mapping
     environments = sandbox_environments_context_var.get(None)
@@ -71,13 +76,19 @@ async def sandbox_with(file: str, on_path: bool = False) -> SandboxEnvironment |
         raise_no_sandbox()
 
     # if we've already discovered the sandbox for this file then return it
-    environment_with_key = f"{file}:{on_path}"
+    environment_with_key = f"{name or ''}:{file}:{on_path}"
     environment = environments_with.get(environment_with_key, None)
     if environment is not None:
         return environment
 
-    # look in each sandbox
-    for _, environment in environments.items():
+    # look in each (or the named) sandbox
+    for environment in (
+        environments.values()
+        if name is None
+        else [named_env]
+        if (named_env := environments.get(name, None))
+        else []
+    ):
         try:
             if on_path:
                 # can we find the file on the path?

inspect_ai/util/_sandbox/docker/compose.py

@@ -96,7 +96,7 @@ async def compose_cp(
     output_limit: int | None = None,
 ) -> None:
     result = await compose_command(
-        ["cp", "--", src, dest],
+        ["cp", "-L", "--", src, dest],
         project=project,
         timeout=120,  # 2-minute timeout for file copies
         cwd=cwd,